




'''simple docstring'''







import argparse
import json
from pathlib import Path

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

from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
_lowerCamelCase				:		Tuple				      =    logging.get_logger(__name__)






def 						_lowercase						(      lowercase_    ):



		'''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      			=			21841
				_snake_case      			=			"imagenet-22k-id2label.json"
				_snake_case      			=			(1, 21841)

		_snake_case      			=			"huggingface/label-files"
		_snake_case      			=			num_labels
		_snake_case      			=			json.load(open(hf_hub_download(lowercase_		,				lowercase_		,				repo_type="dataset"    )		,				"r"    )    )
		_snake_case      			=			{int(lowercase_    ): 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[15027]
		_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 						_lowercase						(      lowercase_    ):



		'''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 						_lowercase						(      ):



		'''simple docstring'''







		_snake_case      			=			"http://images.cocodataset.org/val2017/000000039769.jpg"
		_snake_case      			=			Image.open(requests.get(lowercase_		,				stream=lowercase_    ).raw    )
		return im






@torch.no_grad()
def 						_lowercase						(      lowercase_		,				lowercase_    ):



		'''simple docstring'''







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

		# load HuggingFace model
		_snake_case      			=			ConvNextForImageClassification(lowercase_    )
		model.load_state_dict(lowercase_    )
		model.eval()

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

		_snake_case      			=			model(lowercase_    ).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.1210, -0.6605, 0.1918]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
				_snake_case      			=			torch.tensor([-0.4473, -0.1847, -0.6365]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
				_snake_case      			=			torch.tensor([0.4525, 0.7539, 0.0308]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
				_snake_case      			=			torch.tensor([0.3561, 0.6350, -0.0384]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
				_snake_case      			=			torch.tensor([0.4174, -0.0989, 0.1489]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
				_snake_case      			=			torch.tensor([0.2513, -0.1349, -0.1613]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
				_snake_case      			=			torch.tensor([1.2980, 0.3631, -0.1198]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
				_snake_case      			=			torch.tensor([1.2963, 0.1227, 0.1723]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
				_snake_case      			=			torch.tensor([1.7956, 0.8390, 0.2820]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
				_snake_case      			=			torch.tensor([-0.2822, -0.0502, -0.0878]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
				_snake_case      			=			torch.tensor([-0.5672, -0.0730, -0.4348]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
				_snake_case      			=			torch.tensor([0.2681, 0.2365, 0.6246]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
				_snake_case      			=			torch.tensor([-0.2642, 0.3931, 0.5116]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
				_snake_case      			=			torch.tensor([-0.6677, -0.1873, -0.8379]    )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
				_snake_case      			=			torch.tensor([-0.7749, -0.2967, -0.6444]    )
		else:
				raise ValueError(f'''Unknown URL: {checkpoint_url}'''    )

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

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

		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(lowercase_		,				lowercase_    )		,				organization="nielsr"		,				commit_message="Add model"		,				)


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

						_lowerCamelCase				:		Tuple				      =    parser.parse_args()
						convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)











'''simple docstring'''







import concurrent.futures
import json
import os
import shutil
import tempfile
import unittest

from transformers import AutoTokenizer, PreTrainedTokenizerFast
from transformers.testing_utils import require_tokenizers

from ..test_tokenization_common import TokenizerTesterMixin





@require_tokenizers
class 	lowerCamelCase_				(			UpperCAmelCase_,			unittest.TestCase				):




					'''simple docstring'''






					lowercase__       =					PreTrainedTokenizerFast
					lowercase__       =					False
					lowercase__       =					True
					lowercase__       =					"""tokenizer_file"""






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							_snake_case      			=			False  # because we don't have pretrained_vocab_files_map
							super().setUp()
							_snake_case      			=			True

							_snake_case      			=			["robot-test/dummy-tokenizer-fast", "robot-test/dummy-tokenizer-wordlevel"]
							_snake_case      			=			"SaulLu/dummy-tokenizer-bytelevel-bpe"

							# Inclusion of 2 tokenizers to test different types of models (Unigram and WordLevel for the moment)
							_snake_case      			=			[(PreTrainedTokenizerFast, model_path, {}) for model_path in model_paths]

							_snake_case      			=			PreTrainedTokenizerFast.from_pretrained(model_paths[0]					)
							tokenizer.save_pretrained(self.tmpdirname					)






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							# We disable this test for PreTrainedTokenizerFast because it is the only tokenizer that is not linked to any
							# model
							pass






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							# We disable this test for PreTrainedTokenizerFast because it is the only tokenizer that is not linked to any
							# model
							pass






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							# We disable this test for PreTrainedTokenizerFast because it is the only tokenizer that is not linked to any
							# model
							pass






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							# PreTrainedTokenizerFast doesn't have tokenizer_file in its signature
							pass






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							_snake_case      			=			self.tmpdirname
							# Here we want to test the 2 available tokenizers that use 2 different types of models: Unigram and WordLevel.
							for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
									with self.subTest(f'''{tokenizer.__class__.__name__} ({pretrained_name})'''					):
											try:
													_snake_case      			=			tempfile.mkdtemp()
													_snake_case      			=			self.rust_tokenizer_class.from_pretrained(__UpperCamelCase,    **__UpperCamelCase					)

													tokenizer.save_pretrained(self.tmpdirname					)
													super().test_training_new_tokenizer()
											finally:
													# Even if the test fails, we must be sure that the folder is deleted and that the default tokenizer
													# is restored
													shutil.rmtree(self.tmpdirname					)
													_snake_case      			=			tmpdirname_orig






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							_snake_case      			=			self.tmpdirname
							# Here we want to test the 2 available tokenizers that use 2 different types of models: Unigram and WordLevel.
							for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
									with self.subTest(f'''{tokenizer.__class__.__name__} ({pretrained_name})'''					):
											try:
													_snake_case      			=			tempfile.mkdtemp()
													_snake_case      			=			self.rust_tokenizer_class.from_pretrained(__UpperCamelCase,    **__UpperCamelCase					)

													tokenizer.save_pretrained(self.tmpdirname					)
													super().test_training_new_tokenizer_with_special_tokens_change()
											finally:
													# Even if the test fails, we must be sure that the folder is deleted and that the default tokenizer
													# is restored
													shutil.rmtree(self.tmpdirname					)
													_snake_case      			=			tmpdirname_orig






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							_snake_case      			=			self.rust_tokenizer_class.from_pretrained(self.bytelevel_bpe_model_name					)

							_snake_case      			=			("a" for _ in range(1_000					))
							_snake_case      			=			tokenizer.train_new_from_iterator(text_iterator=__UpperCamelCase,    length=1_000,    vocab_size=50					)

							_snake_case      			=			new_tokenizer.encode("a🤗"					)
							self.assertEqual(__UpperCamelCase,    [64, 172, 253, 97, 245]					)





@require_tokenizers
class 	lowerCamelCase_				(			unittest.TestCase				):




					'''simple docstring'''






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							_snake_case      			=			AutoTokenizer.from_pretrained("bert-base-cased"					)
							_snake_case      			=			json.loads(tokenizer._tokenizer.to_str()					)
							_snake_case      			=			len(__UpperCamelCase					)

							with tempfile.TemporaryDirectory() as tmp_dir:
									# Hack to save this in the tokenizer_config.json
									_snake_case      			=			["tokenizer.4.0.0.json"]
									tokenizer.save_pretrained(__UpperCamelCase					)
									json.dump(__UpperCamelCase,    open(os.path.join(__UpperCamelCase,    "tokenizer.4.0.0.json"					),    "w"					)					)

									# This should pick the new tokenizer file as the version of Transformers is > 4.0.0
									_snake_case      			=			AutoTokenizer.from_pretrained(__UpperCamelCase					)
									self.assertEqual(len(__UpperCamelCase					),    len(__UpperCamelCase					) + 1					)
									_snake_case      			=			json.loads(new_tokenizer._tokenizer.to_str()					)
									self.assertIn("huggingface",    json_tokenizer["model"]["vocab"]					)

									# Will need to be adjusted if we reach v42 and this test is still here.
									# Should pick the old tokenizer file as the version of Transformers is < 4.0.0
									shutil.move(os.path.join(__UpperCamelCase,    "tokenizer.4.0.0.json"					),    os.path.join(__UpperCamelCase,    "tokenizer.42.0.0.json"					)					)
									_snake_case      			=			["tokenizer.42.0.0.json"]
									tokenizer.save_pretrained(__UpperCamelCase					)
									_snake_case      			=			AutoTokenizer.from_pretrained(__UpperCamelCase					)
									self.assertEqual(len(__UpperCamelCase					),    len(__UpperCamelCase					)					)
									_snake_case      			=			json.loads(new_tokenizer._tokenizer.to_str()					)
									self.assertNotIn("huggingface",    json_tokenizer["model"]["vocab"]					)






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							# This repo has two tokenizer files, one for v4.0.0 and above with an added token, one for versions lower.
							_snake_case      			=			"hf-internal-testing/test-two-tokenizers"

							# This should pick the new tokenizer file as the version of Transformers is > 4.0.0
							_snake_case      			=			AutoTokenizer.from_pretrained(__UpperCamelCase					)
							self.assertEqual(len(__UpperCamelCase					),    28_997					)
							_snake_case      			=			json.loads(tokenizer._tokenizer.to_str()					)
							self.assertIn("huggingface",    json_tokenizer["model"]["vocab"]					)

							# Testing an older version by monkey-patching the version in the module it's used.
							import transformers as old_transformers

							_snake_case      			=			"3.0.0"
							_snake_case      			=			old_transformers.models.auto.AutoTokenizer.from_pretrained(__UpperCamelCase					)
							self.assertEqual(len(__UpperCamelCase					),    28_996					)
							_snake_case      			=			json.loads(old_tokenizer._tokenizer.to_str()					)
							self.assertNotIn("huggingface",    json_tokenizer["model"]["vocab"]					)





@require_tokenizers
class 	lowerCamelCase_				(			unittest.TestCase				):




					'''simple docstring'''






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							# See https://github.com/huggingface/transformers/pull/12550
							# and https://github.com/huggingface/tokenizers/issues/537
							_snake_case      			=			PreTrainedTokenizerFast.from_pretrained("robot-test/dummy-tokenizer-wordlevel"					)
							_snake_case      			=			"The Matrix is a 1999 science fiction action film."

							with concurrent.futures.ThreadPoolExecutor() as executor:
									_snake_case      			=			[executor.submit(self.fetch,    __UpperCamelCase,    __UpperCamelCase					) for i in range(10					)]
									_snake_case      			=			[future.result() for future in futures]
									self.assertEqual(__UpperCamelCase,    [[1, 10, 0, 8, 0, 18, 0, 0, 0, 2] for i in range(10					)]					)






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase					):
							return tokenizer.encode(__UpperCamelCase,    truncation="longest_first",    padding="longest"					)





