






"""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()
__A		:	List[str]      =		logging.get_logger(__name__)



def 						lowercase   (    __snake_case    :			Optional[Any]     ):
		lowercase_			:			int								=       ConvNextConfig()

		if "tiny" in checkpoint_url:
				lowercase_			:			Optional[Any]								=       [3, 3, 9, 3]
				lowercase_			:			Union[str, Any]								=       [9_6, 1_9_2, 3_8_4, 7_6_8]
		if "small" in checkpoint_url:
				lowercase_			:			Union[str, Any]								=       [3, 3, 2_7, 3]
				lowercase_			:			Union[str, Any]								=       [9_6, 1_9_2, 3_8_4, 7_6_8]
		if "base" in checkpoint_url:
				lowercase_			:			List[str]								=       [3, 3, 2_7, 3]
				lowercase_			:			int								=       [1_2_8, 2_5_6, 5_1_2, 1_0_2_4]
		if "large" in checkpoint_url:
				lowercase_			:			Union[str, Any]								=       [3, 3, 2_7, 3]
				lowercase_			:			Union[str, Any]								=       [1_9_2, 3_8_4, 7_6_8, 1_5_3_6]
		if "xlarge" in checkpoint_url:
				lowercase_			:			Tuple								=       [3, 3, 2_7, 3]
				lowercase_			:			int								=       [2_5_6, 5_1_2, 1_0_2_4, 2_0_4_8]

		if "1k" in checkpoint_url:
				lowercase_			:			Optional[int]								=       1_0_0_0
				lowercase_			:			Any								=       '''imagenet-1k-id2label.json'''
				lowercase_			:			Union[str, Any]								=       (1, 1_0_0_0)
		else:
				lowercase_			:			Optional[int]								=       2_1_8_4_1
				lowercase_			:			Dict								=       '''imagenet-22k-id2label.json'''
				lowercase_			:			Optional[Any]								=       (1, 2_1_8_4_1)

		lowercase_			:			Union[str, Any]								=       '''huggingface/label-files'''
		lowercase_			:			Optional[Any]								=       num_labels
		lowercase_			:			List[str]								=       json.load(open(hf_hub_download(__snake_case		, __snake_case		, repo_type='''dataset'''     )		, '''r'''     )     )
		lowercase_			:			Any								=       {int(__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[9_2_0_5]
				del idalabel[1_5_0_2_7]
		lowercase_			:			List[Any]								=       idalabel
		lowercase_			:			Any								=       {v: k for k, v in idalabel.items()}
		lowercase_			:			Union[str, Any]								=       hidden_sizes
		lowercase_			:			Dict								=       depths

		return config, expected_shape



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

		return name



def 						lowercase   (    ):
		lowercase_			:			Union[str, Any]								=       '''http://images.cocodataset.org/val2017/000000039769.jpg'''
		lowercase_			:			Union[str, Any]								=       Image.open(requests.get(__snake_case		, stream=__snake_case     ).raw     )
		return im



@torch.no_grad()
def 						lowercase   (    __snake_case    :			List[str]		, __snake_case    :			Union[str, Any]     ):
		lowercase_	, lowercase_			:			str								=       get_convnext_config(__snake_case     )
		# load original state_dict from URL
		lowercase_			:			Tuple								=       torch.hub.load_state_dict_from_url(__snake_case     )['''model''']
		# rename keys
		for key in state_dict.copy().keys():
				lowercase_			:			int								=       state_dict.pop(__snake_case     )
				lowercase_			:			Any								=       val
		# add prefix to all keys expect classifier head
		for key in state_dict.copy().keys():
				lowercase_			:			Union[str, Any]								=       state_dict.pop(__snake_case     )
				if not key.startswith('''classifier'''     ):
						lowercase_			:			Union[str, Any]								=       '''convnext.''' + key
				lowercase_			:			Optional[int]								=       val

		# load HuggingFace model
		lowercase_			:			str								=       ConvNextForImageClassification(__snake_case     )
		model.load_state_dict(__snake_case     )
		model.eval()

		# Check outputs on an image, prepared by ConvNextImageProcessor
		lowercase_			:			Dict								=       2_2_4 if '''224''' in checkpoint_url else 3_8_4
		lowercase_			:			List[Any]								=       ConvNextImageProcessor(size=__snake_case     )
		lowercase_			:			int								=       image_processor(images=prepare_img()		, return_tensors='''pt'''     ).pixel_values

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

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

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

		print('''Pushing model to the hub...'''     )
		lowercase_			:			List[str]								=       '''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(__snake_case		, __snake_case     )		, organization='''nielsr'''		, commit_message='''Add model'''		, )


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

						__A		:	Union[str, Any]      =		parser.parse_args()
						convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)











"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


__A		:	Any      =		{
    '''configuration_realm''': ['''REALM_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''RealmConfig'''],
    '''tokenization_realm''': ['''RealmTokenizer'''],
}

try:
						if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__A		:	Tuple      =		['''RealmTokenizerFast''']

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__A		:	Optional[int]      =		[
						    '''REALM_PRETRAINED_MODEL_ARCHIVE_LIST''',
						    '''RealmEmbedder''',
						    '''RealmForOpenQA''',
						    '''RealmKnowledgeAugEncoder''',
						    '''RealmPreTrainedModel''',
						    '''RealmReader''',
						    '''RealmScorer''',
						    '''load_tf_weights_in_realm''',
						]
						__A		:	Union[str, Any]      =		['''RealmRetriever''']


if TYPE_CHECKING:
						from .configuration_realm import REALM_PRETRAINED_CONFIG_ARCHIVE_MAP, RealmConfig
						from .tokenization_realm import RealmTokenizer

						try:
												if not is_tokenizers_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_realm import RealmTokenizerFast

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_realm import (
												    REALM_PRETRAINED_MODEL_ARCHIVE_LIST,
												    RealmEmbedder,
												    RealmForOpenQA,
												    RealmKnowledgeAugEncoder,
												    RealmPreTrainedModel,
												    RealmReader,
												    RealmScorer,
												    load_tf_weights_in_realm,
												)
												from .retrieval_realm import RealmRetriever


else:
						import sys

						__A		:	Optional[int]      =		_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



