



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__	:  str 		=logging.get_logger(__name__)






def 			__lowercase     (				a__     )       ->				Optional[int]:
							__SCREAMING_SNAKE_CASE								=      ConvNextConfig()

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

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

							__SCREAMING_SNAKE_CASE								=      'huggingface/label-files'
							__SCREAMING_SNAKE_CASE								=      num_labels
							__SCREAMING_SNAKE_CASE								=      json.load(open(hf_hub_download(a__							,     a__							,     repo_type='dataset'     )							,     'r'     )     )
							__SCREAMING_SNAKE_CASE								=      {int(a__     ): 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]
							__SCREAMING_SNAKE_CASE								=      idalabel
							__SCREAMING_SNAKE_CASE								=      {v: k for k, v in idalabel.items()}
							__SCREAMING_SNAKE_CASE								=      hidden_sizes
							__SCREAMING_SNAKE_CASE								=      depths

							return config, expected_shape






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

							return name






def 			__lowercase     (				)       ->				int:
							__SCREAMING_SNAKE_CASE								=      'http://images.cocodataset.org/val2017/000000039769.jpg'
							__SCREAMING_SNAKE_CASE								=      Image.open(requests.get(a__							,     stream=a__     ).raw     )
							return im






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

							# load HuggingFace model
							__SCREAMING_SNAKE_CASE								=      ConvNextForImageClassification(a__     )
							model.load_state_dict(a__     )
							model.eval()

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

							__SCREAMING_SNAKE_CASE								=      model(a__     ).logits

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

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

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

							print('Pushing model to the hub...'     )
							__SCREAMING_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(a__							,     a__     )							,     organization='nielsr'							,     commit_message='Add model'							,     )


if __name__ == "__main__":
			lowerCAmelCase__	:  List[str] 		=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__	:  List[Any] 		=parser.parse_args()
			convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)











import unittest

from cer import CER


lowerCAmelCase__	:  Dict 		=CER()

class      UpperCAmelCase_						(     unittest.TestCase ):

						'''simple docstring'''




						def  _A			(					self	):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE								=      ['White House']
													__SCREAMING_SNAKE_CASE								=      ['white house']
													# S = 2, D = 0, I = 0, N = 11, CER = 2 / 11
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertTrue(abs(char_error_rate - 0.1_8_1_8_1_8_1_8_1_8	) < 1e-6	)






						def  _A			(					self	):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE								=      ['were wolf']
													__SCREAMING_SNAKE_CASE								=      ['werewolf']
													# S = 0, D = 0, I = 1, N = 9, CER = 1 / 9
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertTrue(abs(char_error_rate - 0.1_1_1_1_1_1_1	) < 1e-6	)

													__SCREAMING_SNAKE_CASE								=      ['werewolf']
													__SCREAMING_SNAKE_CASE								=      ['weae     wolf']
													# S = 1, D = 1, I = 0, N = 8, CER = 0.25
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertTrue(abs(char_error_rate - 0.2_5	) < 1e-6	)

													# consecutive whitespaces case 1
													__SCREAMING_SNAKE_CASE								=      ['were wolf']
													__SCREAMING_SNAKE_CASE								=      ['were               wolf']
													# S = 0, D = 0, I = 0, N = 9, CER = 0
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertTrue(abs(char_error_rate - 0.0	) < 1e-6	)

													# consecutive whitespaces case 2
													__SCREAMING_SNAKE_CASE								=      ['were   wolf']
													__SCREAMING_SNAKE_CASE								=      ['were               wolf']
													# S = 0, D = 0, I = 0, N = 9, CER = 0
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertTrue(abs(char_error_rate - 0.0	) < 1e-6	)






						def  _A			(					self	):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE								=      ['werewolf']
													__SCREAMING_SNAKE_CASE								=      ['weaewolf']
													# S = 1, D = 0, I = 0, N = 8, CER = 0.125
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertTrue(abs(char_error_rate - 0.1_2_5	) < 1e-6	)






						def  _A			(					self	):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE								=      ['werewolf']
													__SCREAMING_SNAKE_CASE								=      ['wereawolf']
													# S = 0, D = 1, I = 0, N = 8, CER = 0.125
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertTrue(abs(char_error_rate - 0.1_2_5	) < 1e-6	)






						def  _A			(					self	):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE								=      ['werewolf']
													__SCREAMING_SNAKE_CASE								=      ['wereolf']
													# S = 0, D = 0, I = 1, N = 8, CER = 0.125
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertTrue(abs(char_error_rate - 0.1_2_5	) < 1e-6	)






						def  _A			(					self	):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE								=      ['werewolf']
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertEqual(_A   ,							0.0	)






						def  _A			(					self	):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE								=      ['werewolf', 'I am your father']
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertEqual(_A   ,							0.0	)

													__SCREAMING_SNAKE_CASE								=      ['werewolf', 'I am your father', 'doge']
													__SCREAMING_SNAKE_CASE								=      ['werxwolf', 'I       am your father', 'doge']
													# S = 1, D = 0, I = 0, N = 28, CER = 1 / 28
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertTrue(abs(char_error_rate - 0.0_3_5_7_1_4_2_8	) < 1e-6	)






						def  _A			(					self	):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE								=      ['My hovercraft', 'is full of eels']
													__SCREAMING_SNAKE_CASE								=      ['My hovercraft is full', ' of eels']
													# S = 0, D = 0, I = 2, N = 28, CER = 2 / 28
													# whitespace at the front of " of eels" will be strip during preporcessing
													# so need to insert 2 whitespaces
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A   ,							concatenate_texts=_A	)
													self.assertTrue(abs(char_error_rate - 0.0_7_1_4_2_8	) < 1e-6	)






						def  _A			(					self	):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE								=      ['我能吞下玻璃而不伤身体']
													__SCREAMING_SNAKE_CASE								=      [' 能吞虾玻璃而 不霜身体啦']
													# S = 3, D = 2, I = 0, N = 11, CER = 5 / 11
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertTrue(abs(char_error_rate - 0.4_5_4_5_4_5_4_5_4_5	) < 1e-6	)

													__SCREAMING_SNAKE_CASE								=      ['我能吞下玻璃', '而不伤身体']
													__SCREAMING_SNAKE_CASE								=      ['我    能 吞 下 玻 璃', '而不伤身体']
													# S = 0, D = 5, I = 0, N = 11, CER = 5 / 11
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertTrue(abs(char_error_rate - 0.4_5_4_5_4_5_4_5_4_5_4_5	) < 1e-6	)

													__SCREAMING_SNAKE_CASE								=      ['我能吞下玻璃而不伤身体']
													__SCREAMING_SNAKE_CASE								=      cer.compute(predictions=_A   ,							references=_A	)
													self.assertFalse(_A   ,							0.0	)






						def  _A			(					self	):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE								=      ['']
													__SCREAMING_SNAKE_CASE								=      ['Hypothesis']
													with self.assertRaises(_A	):
																				cer.compute(predictions=_A   ,							references=_A	)


if __name__ == "__main__":
			unittest.main()






