






import inspect
import unittest

from transformers import ConvNextConfig
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_backbone_common import BackboneTesterMixin
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 ConvNextBackbone, ConvNextForImageClassification, ConvNextModel
							from transformers.models.convnext.modeling_convnext import CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							from PIL import Image

							from transformers import AutoImageProcessor
class       _UpperCAmelCase					:

				"""simple docstring"""




				def __init__(     self						:    Optional[int],  lowerCamelCase						:    List[Any],  lowerCamelCase						:    List[Any]=13,  lowerCamelCase						:    Union[str, Any]=32,  lowerCamelCase						:    Optional[int]=3,  lowerCamelCase						:    List[str]=4,  lowerCamelCase						:    str=[10, 20, 30, 40],  lowerCamelCase						:    Tuple=[2, 2, 3, 2],  lowerCamelCase						:    List[Any]=True,  lowerCamelCase						:    List[Any]=True,  lowerCamelCase						:    List[str]=37,  lowerCamelCase						:    str="gelu",  lowerCamelCase						:    Tuple=10,  lowerCamelCase						:    int=0.02,  lowerCamelCase						:    List[str]=["stage2", "stage3", "stage4"],  lowerCamelCase						:    List[Any]=[2, 3, 4],  lowerCamelCase						:    int=None,  ):

								'''simple docstring'''
								lowercase__    =			parent
								lowercase__    =			batch_size
								lowercase__    =			image_size
								lowercase__    =			num_channels
								lowercase__    =			num_stages
								lowercase__    =			hidden_sizes
								lowercase__    =			depths
								lowercase__    =			is_training
								lowercase__    =			use_labels
								lowercase__    =			intermediate_size
								lowercase__    =			hidden_act
								lowercase__    =			num_labels
								lowercase__    =			initializer_range
								lowercase__    =			out_features
								lowercase__    =			out_indices
								lowercase__    =			scope




				def 				lowercase__   (     self						:    List[str]  ):

								'''simple docstring'''
								lowercase__    =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

								lowercase__    =			None
								if self.use_labels:
												lowercase__    =			ids_tensor([self.batch_size],  self.num_labels  )

								lowercase__    =			self.get_config()
								return config, pixel_values, labels




				def 				lowercase__   (     self						:    Tuple  ):

								'''simple docstring'''
								return ConvNextConfig(
								    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 				lowercase__   (     self						:    Any,  lowerCamelCase						:    Dict,  lowerCamelCase						:    int,  lowerCamelCase						:    Any  ):

								'''simple docstring'''
								lowercase__    =			ConvNextModel(config=lowerCamelCase  )
								model.to(lowerCamelCase  )
								model.eval()
								lowercase__    =			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 				lowercase__   (     self						:    Any,  lowerCamelCase						:    Dict,  lowerCamelCase						:    str,  lowerCamelCase						:    List[Any]  ):

								'''simple docstring'''
								lowercase__    =			ConvNextForImageClassification(lowerCamelCase  )
								model.to(lowerCamelCase  )
								model.eval()
								lowercase__    =			model(lowerCamelCase,  labels=lowerCamelCase  )
								self.parent.assertEqual(result.logits.shape,  (self.batch_size, self.num_labels)  )




				def 				lowercase__   (     self						:    Optional[int],  lowerCamelCase						:    int,  lowerCamelCase						:    int,  lowerCamelCase						:    Tuple  ):

								'''simple docstring'''
								lowercase__    =			ConvNextBackbone(config=lowerCamelCase  )
								model.to(lowerCamelCase  )
								model.eval()
								lowercase__    =			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
								lowercase__    =			None
								lowercase__    =			ConvNextBackbone(config=lowerCamelCase  )
								model.to(lowerCamelCase  )
								model.eval()
								lowercase__    =			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 				lowercase__   (     self						:    Optional[int]  ):

								'''simple docstring'''
								lowercase__    =			self.prepare_config_and_inputs()
								lowercase__ ,						lowercase__ ,						lowercase__    =			config_and_inputs
								lowercase__    =			{'''pixel_values''': pixel_values}
								return config, inputs_dict



@require_torch
class       _UpperCAmelCase					(       A__		,A__		,unittest.TestCase	):

				"""simple docstring"""




				lowercase__											=       (
				    (
				        ConvNextModel,
				        ConvNextForImageClassification,
				        ConvNextBackbone,
				    )
				    if is_torch_available()
				    else ()
				)
				lowercase__											=       (
				    {"""feature-extraction""": ConvNextModel, """image-classification""": ConvNextForImageClassification}
				    if is_torch_available()
				    else {}
				)

				lowercase__											=       True
				lowercase__											=       False
				lowercase__											=       False
				lowercase__											=       False
				lowercase__											=       False




				def 				lowercase__   (     self						:    Tuple  ):

								'''simple docstring'''
								lowercase__    =			ConvNextModelTester(self  )
								lowercase__    =			ConfigTester(self,  config_class=lowerCamelCase,  has_text_modality=lowerCamelCase,  hidden_size=37  )




				def 				lowercase__   (     self						:    Tuple  ):

								'''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 				lowercase__   (     self						:    str  ):

								'''simple docstring'''
								return




				@unittest.skip(reason='''ConvNext does not use inputs_embeds'''  )
				def 				lowercase__   (     self						:    List[str]  ):

								'''simple docstring'''
								pass




				@unittest.skip(reason='''ConvNext does not support input and output embeddings'''  )
				def 				lowercase__   (     self						:    Tuple  ):

								'''simple docstring'''
								pass




				@unittest.skip(reason='''ConvNext does not use feedforward chunking'''  )
				def 				lowercase__   (     self						:    Any  ):

								'''simple docstring'''
								pass




				def 				lowercase__   (     self						:    str  ):

								'''simple docstring'''
								lowercase__ ,						lowercase__    =			self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												lowercase__    =			model_class(lowerCamelCase  )
												lowercase__    =			inspect.signature(model.forward  )
												# signature.parameters is an OrderedDict => so arg_names order is deterministic
												lowercase__    =			[*signature.parameters.keys()]

												lowercase__    =			['''pixel_values''']
												self.assertListEqual(arg_names[:1],  lowerCamelCase  )




				def 				lowercase__   (     self						:    Tuple  ):

								'''simple docstring'''
								lowercase__    =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowerCamelCase  )




				def 				lowercase__   (     self						:    Tuple  ):

								'''simple docstring'''
								lowercase__    =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_backbone(*lowerCamelCase  )




				def 				lowercase__   (     self						:    int  ):

								'''simple docstring'''
								def check_hidden_states_output(lowerCamelCase						:    str,  lowerCamelCase						:    Any,  lowerCamelCase						:    Union[str, Any]  ):
												lowercase__    =			model_class(lowerCamelCase  )
												model.to(lowerCamelCase  )
												model.eval()

												with torch.no_grad():
																lowercase__    =			model(**self._prepare_for_class(lowerCamelCase,  lowerCamelCase  )  )

												lowercase__    =			outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

												lowercase__    =			self.model_tester.num_stages
												self.assertEqual(len(lowerCamelCase  ),  expected_num_stages + 1  )

												# ConvNext'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],  )

								lowercase__ ,						lowercase__    =			self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												lowercase__    =			True
												check_hidden_states_output(lowerCamelCase,  lowerCamelCase,  lowerCamelCase  )

												# check that output_hidden_states also work using config
												del inputs_dict["output_hidden_states"]
												lowercase__    =			True

												check_hidden_states_output(lowerCamelCase,  lowerCamelCase,  lowerCamelCase  )




				def 				lowercase__   (     self						:    Any  ):

								'''simple docstring'''
								lowercase__    =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*lowerCamelCase  )




				@slow
				def 				lowercase__   (     self						:    Dict  ):

								'''simple docstring'''
								for model_name in CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												lowercase__    =			ConvNextModel.from_pretrained(lowerCamelCase  )
												self.assertIsNotNone(lowerCamelCase  )






def   a			(						):






				'''simple docstring'''
				lowercase__    =			Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''				)
				return image
@require_torch
@require_vision
class       _UpperCAmelCase					(       unittest.TestCase	):

				"""simple docstring"""




				@cached_property
				def 				lowercase__   (     self						:    Optional[Any]  ):

								'''simple docstring'''
								return AutoImageProcessor.from_pretrained('''facebook/convnext-tiny-224'''  ) if is_vision_available() else None




				@slow
				def 				lowercase__   (     self						:    Dict  ):

								'''simple docstring'''
								lowercase__    =			ConvNextForImageClassification.from_pretrained('''facebook/convnext-tiny-224'''  ).to(lowerCamelCase  )

								lowercase__    =			self.default_image_processor
								lowercase__    =			prepare_img()
								lowercase__    =			image_processor(images=lowerCamelCase,  return_tensors='''pt'''  ).to(lowerCamelCase  )

								# forward pass
								with torch.no_grad():
												lowercase__    =			model(**lowerCamelCase  )

								# verify the logits
								lowercase__    =			torch.Size((1, 1_000)  )
								self.assertEqual(outputs.logits.shape,  lowerCamelCase  )

								lowercase__    =			torch.tensor([-0.0260, -0.4739, 0.1911]  ).to(lowerCamelCase  )

								self.assertTrue(torch.allclose(outputs.logits[0, :3],  lowerCamelCase,  atol=1E-4  )  )



@require_torch
class       _UpperCAmelCase					(       unittest.TestCase		,A__	):

				"""simple docstring"""




				lowercase__											=       (ConvNextBackbone,) if is_torch_available() else ()
				lowercase__											=       ConvNextConfig

				lowercase__											=       False




				def 				lowercase__   (     self						:    Any  ):

								'''simple docstring'''
								lowercase__    =			ConvNextModelTester(self  )












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   a			(						lowerCamelCase_				):






				'''simple docstring'''
				lowercase__    =			[]
				lowercase__    =			[]
				lowercase__    =			[]

				for rt in rc.restypes:
								lowercase__    =			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]				)
								lowercase__    =			{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				)

				lowercase__    =			torch.tensor(
				    lowerCamelCase_		,   dtype=torch.intaa		,   device=protein['''aatype'''].device		,   )
				lowercase__    =			torch.tensor(
				    lowerCamelCase_		,   dtype=torch.intaa		,   device=protein['''aatype'''].device		,   )
				lowercase__    =			torch.tensor(
				    lowerCamelCase_		,   dtype=torch.floataa		,   device=protein['''aatype'''].device		,   )
				lowercase__    =			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
				lowercase__    =			restype_atomaa_to_atomaa[protein_aatype]
				lowercase__    =			restype_atomaa_mask[protein_aatype]

				lowercase__    =			residx_atomaa_mask
				lowercase__    =			residx_atomaa_to_atomaa.long()

				# create the gather indices for mapping back
				lowercase__    =			restype_atomaa_to_atomaa[protein_aatype]
				lowercase__    =			residx_atomaa_to_atomaa.long()

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

				lowercase__    =			restype_atomaa_mask[protein_aatype]
				lowercase__    =			residx_atomaa_mask

				return protein






def   a			(						lowerCamelCase_				):






				'''simple docstring'''
				lowercase__    =			tree_map(lambda lowerCamelCase_				: torch.tensor(lowerCamelCase_		,   device=batch['''aatype'''].device				)		,   lowerCamelCase_		,   np.ndarray				)
				lowercase__    =			tensor_tree_map(lambda lowerCamelCase_				: np.array(lowerCamelCase_				)		,   make_atomaa_masks(lowerCamelCase_				)				)
				return out




