



import inspect
import unittest

from transformers import ConvNextVaConfig
from transformers.models.auto import get_values
from transformers.models.auto.modeling_auto import MODEL_FOR_BACKBONE_MAPPING_NAMES, MODEL_MAPPING_NAMES
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import ConvNextVaBackbone, ConvNextVaForImageClassification, ConvNextVaModel
							from transformers.models.convnextva.modeling_convnextva import CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							from PIL import Image

							from transformers import AutoImageProcessor



class  lowerCamelCase__    :




		'''simple docstring'''



		def __init__(self	,__lowerCamelCase	,__lowerCamelCase=13	,__lowerCamelCase=32	,__lowerCamelCase=3	,__lowerCamelCase=4	,__lowerCamelCase=[10, 20, 30, 40]	,__lowerCamelCase=[2, 2, 3, 2]	,__lowerCamelCase=True	,__lowerCamelCase=True	,__lowerCamelCase=37	,__lowerCamelCase="gelu"	,__lowerCamelCase=10	,__lowerCamelCase=0.02	,__lowerCamelCase=["stage2", "stage3", "stage4"]	,__lowerCamelCase=[2, 3, 4]	,__lowerCamelCase=None	,)  ->    Tuple:

						"""simple docstring"""






						lowerCAmelCase__  :       str				=  parent
						lowerCAmelCase__  :       Optional[Any]				=  batch_size
						lowerCAmelCase__  :       Union[str, Any]				=  image_size
						lowerCAmelCase__  :       Dict				=  num_channels
						lowerCAmelCase__  :       Union[str, Any]				=  num_stages
						lowerCAmelCase__  :       Union[str, Any]				=  hidden_sizes
						lowerCAmelCase__  :       Optional[int]				=  depths
						lowerCAmelCase__  :       Optional[Any]				=  is_training
						lowerCAmelCase__  :       Optional[int]				=  use_labels
						lowerCAmelCase__  :       Tuple				=  intermediate_size
						lowerCAmelCase__  :       List[Any]				=  hidden_act
						lowerCAmelCase__  :       int				=  num_labels
						lowerCAmelCase__  :       List[Any]				=  initializer_range
						lowerCAmelCase__  :       int				=  out_features
						lowerCAmelCase__  :       Optional[Any]				=  out_indices
						lowerCAmelCase__  :       Optional[int]				=  scope



		def 	lowerCAmelCase__   (self				)  ->    Union[str, Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

						lowerCAmelCase__  :       Any				=  None
						if self.use_labels:
										lowerCAmelCase__  :       Union[str, Any]				=  ids_tensor([self.batch_size]	,self.num_labels				)

						lowerCAmelCase__  :       Optional[Any]				=  self.get_config()

						return config, pixel_values, labels



		def 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						return ConvNextVaConfig(
						    num_channels=self.num_channels	,hidden_sizes=self.hidden_sizes	,depths=self.depths	,num_stages=self.num_stages	,hidden_act=self.hidden_act	,is_decoder=__lowerCamelCase	,initializer_range=self.initializer_range	,out_features=self.out_features	,out_indices=self.out_indices	,num_labels=self.num_labels	,)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  ConvNextVaModel(config=__lowerCamelCase				)
						model.to(__lowerCamelCase				)
						model.eval()
						lowerCAmelCase__  :       List[Any]				=  model(__lowerCamelCase				)
						# expected last hidden states: B, C, H // 32, W // 32
						self.parent.assertEqual(
						    result.last_hidden_state.shape	,(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)	,)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       List[Any]				=  ConvNextVaForImageClassification(__lowerCamelCase				)
						model.to(__lowerCamelCase				)
						model.eval()
						lowerCAmelCase__  :       List[str]				=  model(__lowerCamelCase	,labels=__lowerCamelCase				)
						self.parent.assertEqual(result.logits.shape	,(self.batch_size, self.num_labels)				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    Tuple:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  ConvNextVaBackbone(config=__lowerCamelCase				)
						model.to(__lowerCamelCase				)
						model.eval()
						lowerCAmelCase__  :       int				=  model(__lowerCamelCase				)

						# verify hidden states
						self.parent.assertEqual(len(result.feature_maps				)	,len(config.out_features				)				)
						self.parent.assertListEqual(list(result.feature_maps[0].shape				)	,[self.batch_size, self.hidden_sizes[1], 4, 4]				)

						# verify channels
						self.parent.assertEqual(len(model.channels				)	,len(config.out_features				)				)
						self.parent.assertListEqual(model.channels	,config.hidden_sizes[1:]				)

						# verify backbone works with out_features=None
						lowerCAmelCase__  :       Optional[int]				=  None
						lowerCAmelCase__  :       Union[str, Any]				=  ConvNextVaBackbone(config=__lowerCamelCase				)
						model.to(__lowerCamelCase				)
						model.eval()
						lowerCAmelCase__  :       Any				=  model(__lowerCamelCase				)

						# verify feature maps
						self.parent.assertEqual(len(result.feature_maps				)	,1				)
						self.parent.assertListEqual(list(result.feature_maps[0].shape				)	,[self.batch_size, self.hidden_sizes[-1], 1, 1]				)

						# verify channels
						self.parent.assertEqual(len(model.channels				)	,1				)
						self.parent.assertListEqual(model.channels	,[config.hidden_sizes[-1]]				)



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  self.prepare_config_and_inputs()
						lowerCAmelCase__		,	lowerCAmelCase__		,	lowerCAmelCase__  :       int				=  config_and_inputs
						lowerCAmelCase__  :       Tuple				=  {'''pixel_values''': pixel_values}
						return config, inputs_dict



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  self.prepare_config_and_inputs()
						lowerCAmelCase__		,	lowerCAmelCase__		,	lowerCAmelCase__  :       Optional[int]				=  config_and_inputs
						lowerCAmelCase__  :       List[Any]				=  {'''pixel_values''': pixel_values, '''labels''': labels}
						return config, inputs_dict



@require_torch
class  lowerCamelCase__    (							lowerCamelCase__   ,       lowerCamelCase__   ,       unittest.TestCase):




		'''simple docstring'''



		snake_case_			=(
		    (
		        ConvNextVaModel,
		        ConvNextVaForImageClassification,
		        ConvNextVaBackbone,
		    )
		    if is_torch_available()
		    else ()
		)
		snake_case_			=(
		    {"""feature-extraction""": ConvNextVaModel, """image-classification""": ConvNextVaForImageClassification}
		    if is_torch_available()
		    else {}
		)

		snake_case_			=False
		snake_case_			=False
		snake_case_			=False
		snake_case_			=False
		snake_case_			=False



		def 	lowerCAmelCase__   (self				)  ->    Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  ConvNextVaModelTester(self				)
						lowerCAmelCase__  :       Optional[int]				=  ConfigTester(self	,config_class=__lowerCamelCase	,has_text_modality=__lowerCamelCase	,hidden_size=37				)



		def 	lowerCAmelCase__   (self				)  ->    Optional[int]:

						"""simple docstring"""






						self.create_and_test_config_common_properties()
						self.config_tester.create_and_test_config_to_json_string()
						self.config_tester.create_and_test_config_to_json_file()
						self.config_tester.create_and_test_config_from_and_save_pretrained()
						self.config_tester.create_and_test_config_with_num_labels()
						self.config_tester.check_config_can_be_init_without_params()
						self.config_tester.check_config_arguments_init()



		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						return



		@unittest.skip(reason='''ConvNextV2 does not use inputs_embeds'''				)
		def 	lowerCAmelCase__   (self				)  ->    Optional[Any]:

						"""simple docstring"""






						pass



		@unittest.skip(reason='''ConvNextV2 does not support input and output embeddings'''				)
		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						pass



		@unittest.skip(reason='''ConvNextV2 does not use feedforward chunking'''				)
		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						pass



		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						if not self.model_tester.is_training:
										return

						for model_class in self.all_model_classes:
										lowerCAmelCase__		,	lowerCAmelCase__  :       Any				=  self.model_tester.prepare_config_and_inputs_with_labels()
										lowerCAmelCase__  :       Any				=  True

										if model_class.__name__ in [
										    *get_values(__lowerCamelCase				),
										    *get_values(__lowerCamelCase				),
										]:
														continue

										lowerCAmelCase__  :       Any				=  model_class(__lowerCamelCase				)
										model.to(__lowerCamelCase				)
										model.train()
										lowerCAmelCase__  :       Any				=  self._prepare_for_class(__lowerCamelCase	,__lowerCamelCase	,return_labels=__lowerCamelCase				)
										lowerCAmelCase__  :       str				=  model(**__lowerCamelCase				).loss
										loss.backward()



		def 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						if not self.model_tester.is_training:
										return

						for model_class in self.all_model_classes:
										lowerCAmelCase__		,	lowerCAmelCase__  :       List[Any]				=  self.model_tester.prepare_config_and_inputs_with_labels()
										lowerCAmelCase__  :       Tuple				=  False
										lowerCAmelCase__  :       int				=  True

										if (
										    model_class.__name__
										    in [*get_values(__lowerCamelCase				), *get_values(__lowerCamelCase				)]
										    or not model_class.supports_gradient_checkpointing
										):
														continue

										lowerCAmelCase__  :       Dict				=  model_class(__lowerCamelCase				)
										model.to(__lowerCamelCase				)
										model.gradient_checkpointing_enable()
										model.train()
										lowerCAmelCase__  :       Any				=  self._prepare_for_class(__lowerCamelCase	,__lowerCamelCase	,return_labels=__lowerCamelCase				)
										lowerCAmelCase__  :       Optional[int]				=  model(**__lowerCamelCase				).loss
										loss.backward()



		def 	lowerCAmelCase__   (self				)  ->    Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__		,	lowerCAmelCase__  :       Optional[int]				=  self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
										lowerCAmelCase__  :       Union[str, Any]				=  model_class(__lowerCamelCase				)
										lowerCAmelCase__  :       Optional[int]				=  inspect.signature(model.forward				)
										# signature.parameters is an OrderedDict => so arg_names order is deterministic
										lowerCAmelCase__  :       List[str]				=  [*signature.parameters.keys()]

										lowerCAmelCase__  :       Optional[Any]				=  ['''pixel_values''']
										self.assertListEqual(arg_names[:1]	,__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						def check_hidden_states_output(__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				):
										lowerCAmelCase__  :       int				=  model_class(__lowerCamelCase				)
										model.to(__lowerCamelCase				)
										model.eval()

										with torch.no_grad():
														lowerCAmelCase__  :       Dict				=  model(**self._prepare_for_class(__lowerCamelCase	,__lowerCamelCase				)				)

										lowerCAmelCase__  :       List[Any]				=  outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

										lowerCAmelCase__  :       Optional[int]				=  self.model_tester.num_stages
										self.assertEqual(len(__lowerCamelCase				)	,expected_num_stages + 1				)

										# ConvNextV2's feature maps are of shape (batch_size, num_channels, height, width)
										self.assertListEqual(
										    list(hidden_states[0].shape[-2:]				)	,[self.model_tester.image_size // 4, self.model_tester.image_size // 4]	,)

						lowerCAmelCase__		,	lowerCAmelCase__  :       List[str]				=  self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
										lowerCAmelCase__  :       Tuple				=  True
										check_hidden_states_output(__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)

										# check that output_hidden_states also work using config
										del inputs_dict["output_hidden_states"]
										lowerCAmelCase__  :       List[Any]				=  True

										check_hidden_states_output(__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    Union[str, Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*__lowerCamelCase				)



		@slow
		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						for model_name in CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCAmelCase__  :       Dict				=  ConvNextVaModel.from_pretrained(__lowerCamelCase				)
										self.assertIsNotNone(__lowerCamelCase				)






def 			lowerCAmelCase__	(					):





				'''simple docstring'''

				lowerCAmelCase__  :       List[Any]				=  Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''')
				return image



@require_torch
@require_vision
class  lowerCamelCase__    (							unittest.TestCase):




		'''simple docstring'''



		@cached_property
		def 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						return AutoImageProcessor.from_pretrained('''facebook/convnextv2-tiny-1k-224'''				) if is_vision_available() else None



		@slow
		def 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						lowerCAmelCase__  :       str				=  ConvNextVaForImageClassification.from_pretrained('''facebook/convnextv2-tiny-1k-224'''				).to(__lowerCamelCase				)

						lowerCAmelCase__  :       str				=  self.default_image_processor
						lowerCAmelCase__  :       str				=  prepare_img()
						lowerCAmelCase__  :       List[str]				=  preprocessor(images=__lowerCamelCase	,return_tensors='''pt'''				).to(__lowerCamelCase				)

						# forward pass
						with torch.no_grad():
										lowerCAmelCase__  :       List[str]				=  model(**__lowerCamelCase				)

						# verify the logits
						lowerCAmelCase__  :       Any				=  torch.Size((1, 10_00)				)
						self.assertEqual(outputs.logits.shape	,__lowerCamelCase				)

						lowerCAmelCase__  :       str				=  torch.tensor([0.9996, 0.1966, -0.4386]				).to(__lowerCamelCase				)
						self.assertTrue(torch.allclose(outputs.logits[0, :3]	,__lowerCamelCase	,atol=1e-4				)				)






from typing import Dict

import numpy as np
import torch

from . import residue_constants as rc
from .tensor_utils import tensor_tree_map, tree_map






def 			lowerCAmelCase__	(					lowerCamelCase_	:	Dict[str, torch.Tensor]):





				'''simple docstring'''

				lowerCAmelCase__  :       List[Any]				=  []
				lowerCAmelCase__  :       int				=  []
				lowerCAmelCase__  :       Optional[int]				=  []

				for rt in rc.restypes:
								lowerCAmelCase__  :       int				=  rc.restype_name_to_atomaa_names[rc.restype_atoa[rt]]
								restype_atomaa_to_atomaa_list.append([(rc.atom_order[name] if name else 0) for name in atom_names])
								lowerCAmelCase__  :       Any				=  {name: i for i, name in enumerate(lowerCamelCase_)}
								restype_atomaa_to_atomaa_list.append(
								    [(atom_name_to_idxaa[name] if name in atom_name_to_idxaa else 0) for name in rc.atom_types])

								restype_atomaa_mask_list.append([(1.0 if name else 0.0) for name in atom_names])

				# Add dummy mapping for restype 'UNK'
				restype_atomaa_to_atomaa_list.append([0] * 14)
				restype_atomaa_to_atomaa_list.append([0] * 37)
				restype_atomaa_mask_list.append([0.0] * 14)

				lowerCAmelCase__  :       Optional[Any]				=  torch.tensor(
				    lowerCamelCase_ ,dtype=torch.intaa ,device=protein['''aatype'''].device ,)
				lowerCAmelCase__  :       Optional[int]				=  torch.tensor(
				    lowerCamelCase_ ,dtype=torch.intaa ,device=protein['''aatype'''].device ,)
				lowerCAmelCase__  :       List[str]				=  torch.tensor(
				    lowerCamelCase_ ,dtype=torch.floataa ,device=protein['''aatype'''].device ,)
				lowerCAmelCase__  :       Optional[Any]				=  protein['''aatype'''].to(torch.long)

				# create the mapping for (residx, atom14) --> atom37, i.e. an array
				# with shape (num_res, 14) containing the atom37 indices for this protein
				lowerCAmelCase__  :       Optional[int]				=  restype_atomaa_to_atomaa[protein_aatype]
				lowerCAmelCase__  :       List[str]				=  restype_atomaa_mask[protein_aatype]

				lowerCAmelCase__  :       Tuple				=  residx_atomaa_mask
				lowerCAmelCase__  :       Dict				=  residx_atomaa_to_atomaa.long()

				# create the gather indices for mapping back
				lowerCAmelCase__  :       Union[str, Any]				=  restype_atomaa_to_atomaa[protein_aatype]
				lowerCAmelCase__  :       int				=  residx_atomaa_to_atomaa.long()

				# create the corresponding mask
				lowerCAmelCase__  :       List[str]				=  torch.zeros([21, 37] ,dtype=torch.floataa ,device=protein['''aatype'''].device)
				for restype, restype_letter in enumerate(rc.restypes):
								lowerCAmelCase__  :       List[str]				=  rc.restype_atoa[restype_letter]
								lowerCAmelCase__  :       List[str]				=  rc.residue_atoms[restype_name]
								for atom_name in atom_names:
												lowerCAmelCase__  :       List[Any]				=  rc.atom_order[atom_name]
												lowerCAmelCase__  :       Any				=  1

				lowerCAmelCase__  :       int				=  restype_atomaa_mask[protein_aatype]
				lowerCAmelCase__  :       int				=  residx_atomaa_mask

				return protein






def 			lowerCAmelCase__	(					lowerCamelCase_	:	Dict[str, torch.Tensor]):





				'''simple docstring'''

				lowerCAmelCase__  :       Optional[Any]				=  tree_map(lambda lowerCamelCase_: torch.tensor(lowerCamelCase_ ,device=batch['''aatype'''].device) ,lowerCamelCase_ ,np.ndarray)
				lowerCAmelCase__  :       Union[str, Any]				=  tensor_tree_map(lambda lowerCamelCase_: np.array(lowerCamelCase_) ,make_atomaa_masks(lowerCamelCase_))
				return out

