



"""simple docstring"""

import unittest

from transformers import SpeechaTextaConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers.models.speech_to_text_a.modeling_speech_to_text_a import (
							    SpeechaTextaDecoder,
							    SpeechaTextaForCausalLM,
							)


@require_torch
class 			_snake_case     :







						def __init__(self     ,							UpperCAmelCase_     ,							UpperCAmelCase_=9_9     ,							UpperCAmelCase_=1_3     ,							UpperCAmelCase_=1_6     ,							UpperCAmelCase_=7     ,							UpperCAmelCase_=True     ,							UpperCAmelCase_=True     ,							UpperCAmelCase_=True     ,							UpperCAmelCase_=False     ,							UpperCAmelCase_=True     ,							UpperCAmelCase_=2     ,							UpperCAmelCase_=3_2     ,							UpperCAmelCase_=4     ,							UpperCAmelCase_=4     ,							UpperCAmelCase_=3_0     ,							UpperCAmelCase_=0     ,							UpperCAmelCase_=1     ,							UpperCAmelCase_=2     ,							UpperCAmelCase_=None     ,							)			->						Any:
											__lowercase												=			parent
											__lowercase												=			batch_size
											__lowercase												=			decoder_seq_length
											# For common tests
											__lowercase												=			self.decoder_seq_length
											__lowercase												=			is_training
											__lowercase												=			use_attention_mask
											__lowercase												=			use_labels

											__lowercase												=			vocab_size
											__lowercase												=			d_model
											__lowercase												=			d_model
											__lowercase												=			decoder_layers
											__lowercase												=			decoder_layers
											__lowercase												=			decoder_ffn_dim
											__lowercase												=			decoder_attention_heads
											__lowercase												=			decoder_attention_heads
											__lowercase												=			eos_token_id
											__lowercase												=			bos_token_id
											__lowercase												=			pad_token_id
											__lowercase												=			decoder_start_token_id
											__lowercase												=			use_cache
											__lowercase												=			max_position_embeddings

											__lowercase												=			None
											__lowercase												=			decoder_seq_length
											__lowercase												=			2
											__lowercase												=			1







						def        _UpperCamelCase    (self							)			->						Optional[Any]:
											__lowercase												=			ids_tensor([self.batch_size, self.decoder_seq_length]     ,							self.vocab_size							)

											__lowercase												=			None
											if self.use_attention_mask:
																__lowercase												=			ids_tensor([self.batch_size, self.decoder_seq_length]     ,							vocab_size=2							)

											__lowercase												=			None
											if self.use_labels:
																__lowercase												=			ids_tensor([self.batch_size, self.decoder_seq_length]     ,							self.vocab_size							)

											__lowercase												=			SpeechaTextaConfig(
											    vocab_size=self.vocab_size     ,							d_model=self.d_model     ,							decoder_layers=self.decoder_layers     ,							decoder_ffn_dim=self.decoder_ffn_dim     ,							decoder_attention_heads=self.decoder_attention_heads     ,							eos_token_id=self.eos_token_id     ,							bos_token_id=self.bos_token_id     ,							use_cache=self.use_cache     ,							pad_token_id=self.pad_token_id     ,							decoder_start_token_id=self.decoder_start_token_id     ,							max_position_embeddings=self.max_position_embeddings     ,							)

											return (
											    config,
											    input_ids,
											    attention_mask,
											    lm_labels,
											)







						def        _UpperCamelCase    (self     ,							UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_     ,							)			->						List[Any]:
											__lowercase												=			True
											__lowercase												=			SpeechaTextaDecoder(config=UpperCAmelCase_							).to(UpperCAmelCase_							).eval()
											__lowercase												=			input_ids[:2]

											input_ids[input_ids == 0] += 1
											# first forward pass
											__lowercase												=			model(UpperCAmelCase_     ,							use_cache=UpperCAmelCase_							)
											__lowercase												=			model(UpperCAmelCase_							)
											__lowercase												=			model(UpperCAmelCase_     ,							use_cache=UpperCAmelCase_							)

											self.parent.assertTrue(len(UpperCAmelCase_							) == len(UpperCAmelCase_							)							)
											self.parent.assertTrue(len(UpperCAmelCase_							) == len(UpperCAmelCase_							) + 1							)

											__lowercase												=			outputs['''past_key_values''']

											# create hypothetical next token and extent to next_input_ids
											__lowercase												=			ids_tensor((2, 1)     ,							config.vocab_size - 1							) + 1

											# append to next input_ids and
											__lowercase												=			torch.cat([input_ids, next_tokens]     ,							dim=-1							)

											print(UpperCAmelCase_							)

											__lowercase												=			model(UpperCAmelCase_							)['''last_hidden_state''']
											__lowercase												=			model(UpperCAmelCase_     ,							past_key_values=UpperCAmelCase_							)['''last_hidden_state''']

											# select random slice
											__lowercase												=			ids_tensor((1,)     ,							output_from_past.shape[-1]							).item()
											__lowercase												=			output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
											__lowercase												=			output_from_past[:, 0, random_slice_idx].detach()

											# test that outputs are equal for slice
											assert torch.allclose(UpperCAmelCase_     ,							UpperCAmelCase_     ,							atol=1e-3							)







						def        _UpperCamelCase    (self							)			->						int:
											__lowercase												=			self.prepare_config_and_inputs()
											(
											    (
											    __lowercase 
											)						, (
											    __lowercase 
											)						, (
											    __lowercase 
											)						, (
											    __lowercase 
											)						,  
											)												=			config_and_inputs

											__lowercase												=			{
											    '''input_ids''': input_ids,
											    '''attention_mask''': attention_mask,
											}
											return config, inputs_dict



@require_torch
class 			_snake_case     (						UpperCamelCase						,			UpperCamelCase						,			UpperCamelCase						,			unittest.TestCase       ):
						__snake_case 			=     (SpeechaTextaDecoder, SpeechaTextaForCausalLM) if is_torch_available() else ()
						__snake_case 			=     (SpeechaTextaForCausalLM,) if is_torch_available() else ()
						__snake_case 			=     {'text-generation': SpeechaTextaForCausalLM} if is_torch_available() else {}
						__snake_case 			=     True
						__snake_case 			=     False







						def        _UpperCamelCase    (self     ,							)			->						Any:
											__lowercase												=			SpeechaTextaStandaloneDecoderModelTester(self     ,							is_training=UpperCAmelCase_							)
											__lowercase												=			ConfigTester(self     ,							config_class=UpperCAmelCase_							)







						def        _UpperCamelCase    (self							)			->						List[Any]:
											pass







						def        _UpperCamelCase    (self							)			->						Tuple:
											pass







						def        _UpperCamelCase    (self							)			->						Dict:
											pass







						def        _UpperCamelCase    (self							)			->						Tuple:
											self.config_tester.run_common_tests()







						def        _UpperCamelCase    (self							)			->						str:
											__lowercase												=			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_decoder_model_past(*UpperCAmelCase_							)







						def        _UpperCamelCase    (self							)			->						Optional[Any]:
											return





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





def 				__UpperCamelCase   (    a__: str			)				->			Optional[int]:




					"""simple docstring"""



					__lowercase												=			ConvNextConfig()

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

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

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

					return config, expected_shape





def 				__UpperCamelCase   (    a__: Tuple			)				->			List[Any]:




					"""simple docstring"""



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

					return name





def 				__UpperCamelCase   (    )				->			Optional[int]:




					"""simple docstring"""



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





@torch.no_grad()
def 				__UpperCamelCase   (    a__: Union[str, Any] ,	a__: Optional[Any]			)				->			int:




					"""simple docstring"""



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

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

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

					__lowercase												=			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":
										__lowercase												=			torch.tensor([-0.12_10, -0.66_05, 0.19_18]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
										__lowercase												=			torch.tensor([-0.44_73, -0.18_47, -0.63_65]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
										__lowercase												=			torch.tensor([0.45_25, 0.75_39, 0.03_08]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
										__lowercase												=			torch.tensor([0.35_61, 0.63_50, -0.03_84]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
										__lowercase												=			torch.tensor([0.41_74, -0.09_89, 0.14_89]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
										__lowercase												=			torch.tensor([0.25_13, -0.13_49, -0.16_13]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
										__lowercase												=			torch.tensor([1.29_80, 0.36_31, -0.11_98]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
										__lowercase												=			torch.tensor([1.29_63, 0.12_27, 0.17_23]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
										__lowercase												=			torch.tensor([1.79_56, 0.83_90, 0.28_20]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
										__lowercase												=			torch.tensor([-0.28_22, -0.05_02, -0.08_78]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
										__lowercase												=			torch.tensor([-0.56_72, -0.07_30, -0.43_48]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
										__lowercase												=			torch.tensor([0.26_81, 0.23_65, 0.62_46]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
										__lowercase												=			torch.tensor([-0.26_42, 0.39_31, 0.51_16]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
										__lowercase												=			torch.tensor([-0.66_77, -0.18_73, -0.83_79]			)
					elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
										__lowercase												=			torch.tensor([-0.77_49, -0.29_67, -0.64_44]			)
					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...'''			)
					__lowercase												=			'''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__":
							a__ :	Optional[int]						= 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__ :	Any						= parser.parse_args()
							convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)


