



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()
UpperCAmelCase				:							Optional[int]								=logging.get_logger(__name__)



def 				_lowerCAmelCase      (_lowerCAmelCase):
					UpperCamelCase_											=				ConvNextConfig()

					if "tiny" in checkpoint_url:
										UpperCamelCase_											=				[3, 3, 9, 3]
										UpperCamelCase_											=				[96, 1_92, 3_84, 7_68]
					if "small" in checkpoint_url:
										UpperCamelCase_											=				[3, 3, 27, 3]
										UpperCamelCase_											=				[96, 1_92, 3_84, 7_68]
					if "base" in checkpoint_url:
										UpperCamelCase_											=				[3, 3, 27, 3]
										UpperCamelCase_											=				[1_28, 2_56, 5_12, 10_24]
					if "large" in checkpoint_url:
										UpperCamelCase_											=				[3, 3, 27, 3]
										UpperCamelCase_											=				[1_92, 3_84, 7_68, 15_36]
					if "xlarge" in checkpoint_url:
										UpperCamelCase_											=				[3, 3, 27, 3]
										UpperCamelCase_											=				[2_56, 5_12, 10_24, 20_48]

					if "1k" in checkpoint_url:
										UpperCamelCase_											=				10_00
										UpperCamelCase_											=				"imagenet-1k-id2label.json"
										UpperCamelCase_											=				(1, 10_00)
					else:
										UpperCamelCase_											=				2_18_41
										UpperCamelCase_											=				"imagenet-22k-id2label.json"
										UpperCamelCase_											=				(1, 2_18_41)

					UpperCamelCase_											=				"huggingface/label-files"
					UpperCamelCase_											=				num_labels
					UpperCamelCase_											=				json.load(open(hf_hub_download(_lowerCAmelCase				,		_lowerCAmelCase				,		repo_type="dataset")				,		"r"))
					UpperCamelCase_											=				{int(_lowerCAmelCase): 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[92_05]
										del idalabel[1_50_27]
					UpperCamelCase_											=				idalabel
					UpperCamelCase_											=				{v: k for k, v in idalabel.items()}
					UpperCamelCase_											=				hidden_sizes
					UpperCamelCase_											=				depths

					return config, expected_shape



def 				_lowerCAmelCase      (_lowerCAmelCase):
					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 "stages" in name:
										UpperCamelCase_											=				name.replace("stages"				,		"encoder.stages")
					if "norm" in name:
										UpperCamelCase_											=				name.replace("norm"				,		"layernorm")
					if "gamma" in name:
										UpperCamelCase_											=				name.replace("gamma"				,		"layer_scale_parameter")
					if "head" in name:
										UpperCamelCase_											=				name.replace("head"				,		"classifier")

					return name



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



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

					# load HuggingFace model
					UpperCamelCase_											=				ConvNextForImageClassification(_lowerCAmelCase)
					model.load_state_dict(_lowerCAmelCase)
					model.eval()

					# Check outputs on an image, prepared by ConvNextImageProcessor
					UpperCamelCase_											=				2_24 if "224" in checkpoint_url else 3_84
					UpperCamelCase_											=				ConvNextImageProcessor(size=_lowerCAmelCase)
					UpperCamelCase_											=				image_processor(images=prepare_img()				,		return_tensors="pt").pixel_values

					UpperCamelCase_											=				model(_lowerCAmelCase).logits

					# note: the logits below were obtained without center cropping
					if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
										UpperCamelCase_											=				torch.tensor([-0.1210, -0.6605, 0.1918])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
										UpperCamelCase_											=				torch.tensor([-0.4473, -0.1847, -0.6365])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
										UpperCamelCase_											=				torch.tensor([0.4525, 0.7539, 0.0308])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
										UpperCamelCase_											=				torch.tensor([0.3561, 0.6350, -0.0384])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
										UpperCamelCase_											=				torch.tensor([0.4174, -0.0989, 0.1489])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
										UpperCamelCase_											=				torch.tensor([0.2513, -0.1349, -0.1613])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
										UpperCamelCase_											=				torch.tensor([1.2980, 0.3631, -0.1198])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
										UpperCamelCase_											=				torch.tensor([1.2963, 0.1227, 0.1723])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
										UpperCamelCase_											=				torch.tensor([1.7956, 0.8390, 0.2820])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
										UpperCamelCase_											=				torch.tensor([-0.2822, -0.0502, -0.0878])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
										UpperCamelCase_											=				torch.tensor([-0.5672, -0.0730, -0.4348])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
										UpperCamelCase_											=				torch.tensor([0.2681, 0.2365, 0.6246])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
										UpperCamelCase_											=				torch.tensor([-0.2642, 0.3931, 0.5116])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
										UpperCamelCase_											=				torch.tensor([-0.6677, -0.1873, -0.8379])
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
										UpperCamelCase_											=				torch.tensor([-0.7749, -0.2967, -0.6444])
					else:
										raise ValueError(f"""Unknown URL: {checkpoint_url}""")

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

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

					print("Pushing model to the hub...")
					UpperCamelCase_											=				"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(_lowerCAmelCase				,		_lowerCAmelCase)				,		organization="nielsr"				,		commit_message="Add model"				,		)


if __name__ == "__main__":
	UpperCAmelCase				:							Union[str, 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.""",
	)

	UpperCAmelCase				:							str								=parser.parse_args()
	convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)





import json
import os
import shutil
import tempfile
import unittest

from transformers.models.wavaveca import WavaVecaCTCTokenizer, WavaVecaFeatureExtractor, WavaVecaProcessor
from transformers.models.wavaveca.tokenization_wavaveca import VOCAB_FILES_NAMES
from transformers.utils import FEATURE_EXTRACTOR_NAME

from .test_feature_extraction_wavaveca import floats_list


class  _lowercase				(unittest.TestCase ):
				'''simple docstring'''







				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				"<pad> <s> </s> <unk> | E T A O N I H S R D L U M W C F G Y P B V K ' X J Q Z".split(" "      )
									UpperCamelCase_											=				dict(zip(snake_case__       ,	range(len(snake_case__      )      )      )      )

									UpperCamelCase_											=				{
									    "pad_token": "<pad>",
									    "unk_token": "<unk>",
									    "bos_token": "<s>",
									    "eos_token": "</s>",
									}
									UpperCamelCase_											=				{
									    "feature_size": 1,
									    "padding_value": 0.0,
									    "sampling_rate": 1_6000,
									    "return_attention_mask": False,
									    "do_normalize": True,
									}

									UpperCamelCase_											=				tempfile.mkdtemp()
									UpperCamelCase_											=				os.path.join(self.tmpdirname       ,	VOCAB_FILES_NAMES["vocab_file"]      )
									UpperCamelCase_											=				os.path.join(self.tmpdirname       ,	snake_case__      )
									with open(self.vocab_file       ,	"w"       ,	encoding="utf-8"      ) as fp:
														fp.write(json.dumps(snake_case__      ) + "\n"      )

									with open(self.feature_extraction_file       ,	"w"       ,	encoding="utf-8"      ) as fp:
														fp.write(json.dumps(snake_case__      ) + "\n"      )

				def  _lowerCamelCase   (	self       ,	**snake_case__      ):

									'''simple docstring'''


									UpperCamelCase_											=				self.add_kwargs_tokens_map.copy()
									kwargs.update(snake_case__      )
									return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname       ,	**snake_case__      )

				def  _lowerCamelCase   (	self       ,	**snake_case__      ):

									'''simple docstring'''


									return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname       ,	**snake_case__      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									shutil.rmtree(self.tmpdirname      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				self.get_tokenizer()
									UpperCamelCase_											=				self.get_feature_extractor()

									UpperCamelCase_											=				WavaVecaProcessor(tokenizer=snake_case__       ,	feature_extractor=snake_case__      )

									processor.save_pretrained(self.tmpdirname      )
									UpperCamelCase_											=				WavaVecaProcessor.from_pretrained(self.tmpdirname      )

									self.assertEqual(processor.tokenizer.get_vocab()       ,	tokenizer.get_vocab()      )
									self.assertIsInstance(processor.tokenizer       ,	snake_case__      )

									self.assertEqual(processor.feature_extractor.to_json_string()       ,	feature_extractor.to_json_string()      )
									self.assertIsInstance(processor.feature_extractor       ,	snake_case__      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				WavaVecaProcessor(tokenizer=self.get_tokenizer()       ,	feature_extractor=self.get_feature_extractor()      )
									processor.save_pretrained(self.tmpdirname      )

									UpperCamelCase_											=				self.get_tokenizer(bos_token="(BOS)"       ,	eos_token="(EOS)"      )
									UpperCamelCase_											=				self.get_feature_extractor(do_normalize=snake_case__       ,	padding_value=1.0      )

									UpperCamelCase_											=				WavaVecaProcessor.from_pretrained(
									    self.tmpdirname       ,	bos_token="(BOS)"       ,	eos_token="(EOS)"       ,	do_normalize=snake_case__       ,	padding_value=1.0      )

									self.assertEqual(processor.tokenizer.get_vocab()       ,	tokenizer_add_kwargs.get_vocab()      )
									self.assertIsInstance(processor.tokenizer       ,	snake_case__      )

									self.assertEqual(processor.feature_extractor.to_json_string()       ,	feature_extractor_add_kwargs.to_json_string()      )
									self.assertIsInstance(processor.feature_extractor       ,	snake_case__      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				self.get_feature_extractor()
									UpperCamelCase_											=				self.get_tokenizer()

									UpperCamelCase_											=				WavaVecaProcessor(tokenizer=snake_case__       ,	feature_extractor=snake_case__      )

									UpperCamelCase_											=				floats_list((3, 1000)      )

									UpperCamelCase_											=				feature_extractor(snake_case__       ,	return_tensors="np"      )
									UpperCamelCase_											=				processor(snake_case__       ,	return_tensors="np"      )

									for key in input_feat_extract.keys():
														self.assertAlmostEqual(input_feat_extract[key].sum()       ,	input_processor[key].sum()       ,	delta=1e-2      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				self.get_feature_extractor()
									UpperCamelCase_											=				self.get_tokenizer()

									UpperCamelCase_											=				WavaVecaProcessor(tokenizer=snake_case__       ,	feature_extractor=snake_case__      )

									UpperCamelCase_											=				"This is a test string"

									UpperCamelCase_											=				processor(text=snake_case__      )

									UpperCamelCase_											=				tokenizer(snake_case__      )

									for key in encoded_tok.keys():
														self.assertListEqual(encoded_tok[key]       ,	encoded_processor[key]      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				self.get_feature_extractor()
									UpperCamelCase_											=				self.get_tokenizer()

									UpperCamelCase_											=				WavaVecaProcessor(tokenizer=snake_case__       ,	feature_extractor=snake_case__      )

									UpperCamelCase_											=				[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

									UpperCamelCase_											=				processor.batch_decode(snake_case__      )
									UpperCamelCase_											=				tokenizer.batch_decode(snake_case__      )

									self.assertListEqual(snake_case__       ,	snake_case__      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				self.get_feature_extractor()
									UpperCamelCase_											=				self.get_tokenizer()

									UpperCamelCase_											=				WavaVecaProcessor(tokenizer=snake_case__       ,	feature_extractor=snake_case__      )

									self.assertListEqual(
									    processor.model_input_names       ,	feature_extractor.model_input_names       ,	msg="`processor` and `feature_extractor` model input names do not match"       ,	)
