



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





def 						__lowerCamelCase       (      __magic_name__	:  Tuple						):
	a__:       List[Any]											=ConvNextVaConfig()

	if "atto" in checkpoint_url:
		a__:       List[str]											=[2, 2, 6, 2]
		a__:       Optional[Any]											=[40, 80, 160, 320]
	if "femto" in checkpoint_url:
		a__:       List[str]											=[2, 2, 6, 2]
		a__:       Optional[Any]											=[48, 96, 192, 384]
	if "pico" in checkpoint_url:
		a__:       List[Any]											=[2, 2, 6, 2]
		a__:       List[str]											=[64, 128, 256, 512]
	if "nano" in checkpoint_url:
		a__:       List[Any]											=[2, 2, 8, 2]
		a__:       List[str]											=[80, 160, 320, 640]
	if "tiny" in checkpoint_url:
		a__:       List[Any]											=[3, 3, 9, 3]
		a__:       str											=[96, 192, 384, 768]
	if "base" in checkpoint_url:
		a__:       Tuple											=[3, 3, 27, 3]
		a__:       Any											=[128, 256, 512, 1_024]
	if "large" in checkpoint_url:
		a__:       Union[str, Any]											=[3, 3, 27, 3]
		a__:       Optional[int]											=[192, 384, 768, 1_536]
	if "huge" in checkpoint_url:
		a__:       Union[str, Any]											=[3, 3, 27, 3]
		a__:       Optional[int]											=[352, 704, 1_408, 2_816]

	a__:       str											=1_000
	a__:       List[str]											="imagenet-1k-id2label.json"
	a__:       Tuple											=(1, 1_000)

	a__:       int											="huggingface/label-files"
	a__:       str											=num_labels
	a__:       int											=json.load(open(hf_hub_download(__magic_name__     , __magic_name__     , repo_type="dataset"						)     , "r"						)						)
	a__:       Any											={int(__magic_name__						): v for k, v in idalabel.items()}

	a__:       str											=idalabel
	a__:       Dict											={v: k for k, v in idalabel.items()}
	a__:       Any											=hidden_sizes
	a__:       int											=depths

	return config, expected_shape





def 						__lowerCamelCase       (      __magic_name__	:  List[str]						):



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

	return name





def 						__lowerCamelCase       (      ):
	a__:       Optional[Any]											="http://images.cocodataset.org/val2017/000000039769.jpg"
	a__:       Tuple											=Image.open(requests.get(__magic_name__     , stream=__magic_name__						).raw						)
	return im





def 						__lowerCamelCase       (      __magic_name__	:  int						):



	if "224" in checkpoint_url:
		a__:       List[str]											=224
		a__:       Union[str, Any]											=224 / 256
	elif "384" in checkpoint_url:
		a__:       Tuple											=384
		a__:       int											=None
	else:
		a__:       List[str]											=512
		a__:       Any											=None

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





@torch.no_grad()
def 						__lowerCamelCase       (      __magic_name__	:  int     , __magic_name__	:  List[str]     , __magic_name__	:  Optional[Any]     , __magic_name__	:  Optional[int]						):
	print("Downloading original model from checkpoint..."						)
	# define ConvNeXTV2 configuration based on URL
	a__							,			a__:       List[str]											=get_convnextva_config(__magic_name__						)
	# load original state_dict from URL
	a__:       int											=torch.hub.load_state_dict_from_url(__magic_name__						)["model"]

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

	# load HuggingFace model
	a__:       List[str]											=ConvNextVaForImageClassification(__magic_name__						)
	model.load_state_dict(__magic_name__						)
	model.eval()

	# Check outputs on an image, prepared by ConvNextImageProcessor
	a__:       Union[str, Any]											=convert_preprocessor(__magic_name__						)
	a__:       Union[str, Any]											=preprocessor(images=prepare_img()     , return_tensors="pt"						)
	a__:       Union[str, Any]											=model(**__magic_name__						).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":
		a__:       List[Any]											=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":
		a__:       Dict											=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":
		a__:       Tuple											=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":
		a__:       int											=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":
		a__:       str											=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":
		a__:       Optional[int]											=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":
		a__:       Dict											=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":
		a__:       str											=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":
		a__:       Dict											=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":
		a__:       List[str]											=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":
		a__:       Optional[int]											=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":
		a__:       str											=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":
		a__:       List[str]											=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":
		a__:       List[Any]											=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":
		a__:       List[str]											=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":
		a__:       Tuple											=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":
		a__:       int											=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":
		a__:       Tuple											=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]     , __magic_name__     , 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(__magic_name__						):
			os.mkdir(__magic_name__						)

		model.save_pretrained(__magic_name__						)
		preprocessor.save_pretrained(__magic_name__						)

	a__:       int											="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(__magic_name__						)
		preprocessor.push_to_hub(__magic_name__						)


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

		__UpperCAmelCase						  =     parser.parse_args()
		convert_convnextva_checkpoint(
		    args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
		)










import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


__UpperCAmelCase						  =     10_00_00_00_00_00
__UpperCAmelCase						  =     1_00_00

__UpperCAmelCase  ,							__UpperCAmelCase						  =     os.path.split(__file__)
__UpperCAmelCase						  =     os.path.join(RESULTS_BASEPATH, '''results''', RESULTS_FILENAME.replace('''.py''', '''.json'''))





def 						__lowerCamelCase       (      __magic_name__	:  int     , __magic_name__	:  int						):
							a__:       Any											=pa.Table.from_pydict({"col": [0] * chunk_size}						)
							a__:       Optional[Any]											=pa.concat_tables([table] * (num_examples // chunk_size)						)
							return datasets.Dataset(__magic_name__     , fingerprint="table_100B"						)


@dataclass
class 	lowerCamelCase__ :
						_lowerCAmelCase								= 42
						_lowerCAmelCase								= 42
						_lowerCAmelCase								= 42
						_lowerCAmelCase								= 42


						def   _lowerCamelCase			( self						:  Union[str, Any]			):
													a__:       Tuple											=np.random.default_rng(self.seed			)
													a__:       Tuple											=rng.integers(low=self.low      ,       high=self.high      ,       size=self.size			).tolist()


						def __iter__( self						:  Any			):
													return iter(self._sampled_values			)







						def __len__( self						:  Any			):
													return self.size





@get_duration
def 						__lowerCamelCase       (      __magic_name__	:  datasets.Dataset						):
							a__:       Union[str, Any]											=dataset[0]





@get_duration
def 						__lowerCamelCase       (      __magic_name__	:  datasets.Dataset						):
							a__:       Optional[Any]											=dataset[-1]





@get_duration
def 						__lowerCamelCase       (      __magic_name__	:  datasets.Dataset						):
							a__:       List[str]											=dataset[range(len(__magic_name__						) // 2     , len(__magic_name__						) // 2 + 1_024						)]





@get_duration
def 						__lowerCamelCase       (      __magic_name__	:  datasets.Dataset						):
							a__:       Optional[int]											=dataset[RandIter(0     , len(__magic_name__						)     , 1_024     , seed=42						)]





def 						__lowerCamelCase       (      ):
							a__:       Dict											={"num examples": SPEED_TEST_N_EXAMPLES}
							a__:       Union[str, Any]											=(get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
							print("generating dataset"						)
							a__:       Dict											=generate_100B_dataset(num_examples=__magic_name__     , chunk_size=__magic_name__						)
							print("Functions"						)
							for func in functions:
														print(func.__name__						)
														a__:       Dict											=func(__magic_name__						)

							with open(__magic_name__     , "wb"						) as f:
														f.write(json.dumps(__magic_name__						).encode("utf-8"						)						)


if __name__ == "__main__":  # useful to run the profiler
		benchmark_table_100B()





