'''simple docstring'''




from typing import Dict, Iterable, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_DEFAULT_MEAN,
    IMAGENET_DEFAULT_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


__UpperCAmelCase     							=logging.get_logger(__name__)




class     a__      (   UpperCAmelCase__ ):
     lowerCamelCase  :     List[str]							=["pixel_values"]






     def __init__( self		:		List[Any]  ,				a		:		bool = True  ,				a		:		Dict[str, int] = None  ,				a		:		PILImageResampling = PILImageResampling.BICUBIC  ,				a		:		bool = True  ,				a		:		Dict[str, int] = None  ,				a		:		bool = True  ,				a		:		Union[int, float] = 1 / 2_55  ,				a		:		bool = True  ,				a		:		Optional[Union[float, Iterable[float]]] = IMAGENET_DEFAULT_MEAN  ,				a		:		Optional[Union[float, Iterable[float]]] = IMAGENET_DEFAULT_STD  ,				**a		:		List[Any]  ,				):





         """simple docstring"""



         super().__init__(**a    )
         __lowerCamelCase		=  size if size is not None else {'''shortest_edge''': 2_24}
         __lowerCamelCase		=  get_size_dict(a  ,				default_to_square=a    )
         __lowerCamelCase		=  crop_size if crop_size is not None else {'''height''': 2_24, '''width''': 2_24}
         __lowerCamelCase		=  get_size_dict(a  ,				param_name='''crop_size'''    )

         __lowerCamelCase		=  do_resize
         __lowerCamelCase		=  size
         __lowerCamelCase		=  resample
         __lowerCamelCase		=  do_center_crop
         __lowerCamelCase		=  crop_size
         __lowerCamelCase		=  do_rescale
         __lowerCamelCase		=  rescale_factor
         __lowerCamelCase		=  do_normalize
         __lowerCamelCase		=  image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
         __lowerCamelCase		=  image_std if image_std is not None else IMAGENET_DEFAULT_STD






     def    SCREAMING_SNAKE_CASE__   ( self		:		int  ,				a		:		np.ndarray  ,				a		:		Dict[str, int]  ,				a		:		PILImageResampling = PILImageResampling.BICUBIC  ,				a		:		Optional[Union[str, ChannelDimension]] = None  ,				**a		:		List[Any]  ,				):





         """simple docstring"""



         __lowerCamelCase		=  get_size_dict(a  ,				default_to_square=a    )
         # size_dict is a dict with either keys "height" and "width" or "shortest_edge"
         if "shortest_edge" in size:
             __lowerCamelCase		=  int((2_56 / 2_24) * size['''shortest_edge''']    )
             __lowerCamelCase		=  get_resize_output_image_size(a  ,				size=a  ,				default_to_square=a    )
             __lowerCamelCase		=  {'''height''': output_size[0], '''width''': output_size[1]}
         if "height" not in size_dict or "width" not in size_dict:
             raise ValueError(
                 f"""Size dict must have keys 'height' and 'width' or 'shortest_edge'. Got {size_dict.keys()}"""    )
         return resize(
             a  ,				size=(size_dict['''height'''], size_dict['''width'''])  ,				resample=a  ,				data_format=a  ,				**a    )






     def    SCREAMING_SNAKE_CASE__   ( self		:		List[Any]  ,				a		:		np.ndarray  ,				a		:		Dict[str, int]  ,				a		:		Optional[Union[str, ChannelDimension]] = None  ,				**a		:		Tuple  ,				):





         """simple docstring"""



         __lowerCamelCase		=  get_size_dict(a    )
         if "height" not in size or "width" not in size:
             raise ValueError(f"""Size dict must have keys 'height' and 'width'. Got {size.keys()}"""    )
         return center_crop(a  ,				size=(size['''height'''], size['''width'''])  ,				data_format=a  ,				**a    )






     def    SCREAMING_SNAKE_CASE__   ( self		:		Dict  ,				a		:		np.ndarray  ,				a		:		Union[int, float]  ,				a		:		Optional[Union[str, ChannelDimension]] = None  ,				**a		:		Union[str, Any]  ,				):





         """simple docstring"""



         return rescale(a  ,				scale=a  ,				data_format=a  ,				**a    )






     def    SCREAMING_SNAKE_CASE__   ( self		:		int  ,				a		:		np.ndarray  ,				a		:		Union[float, List[float]]  ,				a		:		Union[float, List[float]]  ,				a		:		Optional[Union[str, ChannelDimension]] = None  ,				**a		:		Any  ,				):





         """simple docstring"""



         return normalize(a  ,				mean=a  ,				std=a  ,				data_format=a  ,				**a    )






     def    SCREAMING_SNAKE_CASE__   ( self		:		List[Any]  ,				a		:		ImageInput  ,				a		:		Optional[bool] = None  ,				a		:		Optional[Dict[str, int]] = None  ,				a		:		PILImageResampling = None  ,				a		:		Optional[bool] = None  ,				a		:		Optional[Dict[str, int]] = None  ,				a		:		Optional[bool] = None  ,				a		:		Optional[float] = None  ,				a		:		Optional[bool] = None  ,				a		:		Optional[Union[float, Iterable[float]]] = None  ,				a		:		Optional[Union[float, Iterable[float]]] = None  ,				a		:		Optional[TensorType] = None  ,				a		:		ChannelDimension = ChannelDimension.FIRST  ,				**a		:		str  ,				):





         """simple docstring"""



         __lowerCamelCase		=  do_resize if do_resize is not None else self.do_resize
         __lowerCamelCase		=  resample if resample is not None else self.resample
         __lowerCamelCase		=  do_center_crop if do_center_crop is not None else self.do_center_crop
         __lowerCamelCase		=  do_rescale if do_rescale is not None else self.do_rescale
         __lowerCamelCase		=  rescale_factor if rescale_factor is not None else self.rescale_factor
         __lowerCamelCase		=  do_normalize if do_normalize is not None else self.do_normalize
         __lowerCamelCase		=  image_mean if image_mean is not None else self.image_mean
         __lowerCamelCase		=  image_std if image_std is not None else self.image_std

         __lowerCamelCase		=  size if size is not None else self.size
         __lowerCamelCase		=  get_size_dict(a  ,				default_to_square=a    )
         __lowerCamelCase		=  crop_size if crop_size is not None else self.crop_size
         __lowerCamelCase		=  get_size_dict(a  ,				param_name='''crop_size'''    )

         __lowerCamelCase		=  make_list_of_images(a    )

         if not valid_images(a    ):
             raise ValueError(
                 '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
                 '''torch.Tensor, tf.Tensor or jax.ndarray.'''    )

         if do_resize and size is None:
             raise ValueError('''Size must be specified if do_resize is True.'''    )

         if do_center_crop and crop_size is None:
             raise ValueError('''Crop size must be specified if do_center_crop is True.'''    )

         if do_rescale and rescale_factor is None:
             raise ValueError('''Rescale factor must be specified if do_rescale is True.'''    )

         if do_normalize and (image_mean is None or image_std is None):
             raise ValueError('''Image mean and std must be specified if do_normalize is True.'''    )

         # All transformations expect numpy arrays.
         __lowerCamelCase		=  [to_numpy_array(a    ) for image in images]

         if do_resize:
             __lowerCamelCase		=  [self.resize(a  ,				a  ,				a    ) for image in images]

         if do_center_crop:
             __lowerCamelCase		=  [self.center_crop(a  ,				a    ) for image in images]

         if do_rescale:
             __lowerCamelCase		=  [self.rescale(a  ,				a    ) for image in images]

         if do_normalize:
             __lowerCamelCase		=  [self.normalize(a  ,				a  ,				a    ) for image in images]

         __lowerCamelCase		=  [to_channel_dimension_format(a  ,				a    ) for image in images]

         __lowerCamelCase		=  {'''pixel_values''': images}
         return BatchFeature(data=a  ,				tensor_type=a    )



import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

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


if is_torch_available():
			import torch

			from transformers.models.esm.modeling_esmfold import EsmForProteinFolding

class 				UpperCAmelCase_		:
							"""simple docstring"""







							def __init__(			self       ,							_a       ,							_a=1_3       ,							_a=7       ,							_a=False       ,							_a=True       ,							_a=False       ,							_a=False       ,							_a=1_9       ,							_a=3_2       ,							_a=5       ,							_a=4       ,							_a=3_7       ,							_a="gelu"       ,							_a=0.1       ,							_a=0.1       ,							_a=5_1_2       ,							_a=1_6       ,							_a=2       ,							_a=0.02       ,							_a=3       ,							_a=4       ,							_a=None       ,							)    ->    Union[str, Any]:
												_a	:    Optional[Any]         =		parent
												_a	:    Union[str, Any]         =		batch_size
												_a	:    List[Any]         =		seq_length
												_a	:    Dict         =		is_training
												_a	:    int         =		use_input_mask
												_a	:    str         =		use_token_type_ids
												_a	:    Any         =		use_labels
												_a	:    List[Any]         =		vocab_size
												_a	:    Any         =		hidden_size
												_a	:    int         =		num_hidden_layers
												_a	:    str         =		num_attention_heads
												_a	:    Dict         =		intermediate_size
												_a	:    List[str]         =		hidden_act
												_a	:    Optional[Any]         =		hidden_dropout_prob
												_a	:    Optional[Any]         =		attention_probs_dropout_prob
												_a	:    int         =		max_position_embeddings
												_a	:    Tuple         =		type_vocab_size
												_a	:    str         =		type_sequence_label_size
												_a	:    Any         =		initializer_range
												_a	:    Union[str, Any]         =		num_labels
												_a	:    Dict         =		num_choices
												_a	:    Union[str, Any]         =		scope


							def 				__lowercase				(			self       )    ->    List[Any]:
												_a	:    str         =		ids_tensor([self.batch_size, self.seq_length]       ,							self.vocab_size       )

												_a	:    Dict         =		None
												if self.use_input_mask:
																	_a	:    int         =		random_attention_mask([self.batch_size, self.seq_length]       )

												_a	:    List[Any]         =		None
												_a	:    Tuple         =		None
												_a	:    Any         =		None
												if self.use_labels:
																	_a	:    Tuple         =		ids_tensor([self.batch_size]       ,							self.type_sequence_label_size       )
																	_a	:    Union[str, Any]         =		ids_tensor([self.batch_size, self.seq_length]       ,							self.num_labels       )
																	_a	:    List[Any]         =		ids_tensor([self.batch_size]       ,							self.num_choices       )

												_a	:    str         =		self.get_config()

												return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels


							def 				__lowercase				(			self       )    ->    str:
												_a	:    Optional[int]         =		EsmConfig(
												    vocab_size=3_3       ,							hidden_size=self.hidden_size       ,							pad_token_id=1       ,							num_hidden_layers=self.num_hidden_layers       ,							num_attention_heads=self.num_attention_heads       ,							intermediate_size=self.intermediate_size       ,							hidden_act=self.hidden_act       ,							hidden_dropout_prob=self.hidden_dropout_prob       ,							attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,							max_position_embeddings=self.max_position_embeddings       ,							type_vocab_size=self.type_vocab_size       ,							initializer_range=self.initializer_range       ,							is_folding_model=_a       ,							esmfold_config={'''trunk''': {'''num_blocks''': 2}, '''fp16_esm''': False}       ,							)
												return config


							def 				__lowercase				(			self       ,							_a       ,							_a       ,							_a       ,							_a       ,							_a       ,							_a       )    ->    str:
												_a	:    Union[str, Any]         =		EsmForProteinFolding(config=_a       ).float()
												model.to(_a       )
												model.eval()
												_a	:    str         =		model(_a       ,							attention_mask=_a       )
												_a	:    Union[str, Any]         =		model(_a       )
												_a	:    Optional[int]         =		model(_a       )

												self.parent.assertEqual(result.positions.shape       ,							(8, self.batch_size, self.seq_length, 1_4, 3)       )
												self.parent.assertEqual(result.angles.shape       ,							(8, self.batch_size, self.seq_length, 7, 2)       )




							def 				__lowercase				(			self       )    ->    str:
												_a	:    List[str]         =		self.prepare_config_and_inputs()
												(
												    (
												    _a
												)      ,					(
												    _a
												)      ,					(
												    _a
												)      ,					(
												    _a
												)      ,					(
												    _a
												)      ,					(
												    _a
												)      ,					
												)	:    Optional[Any]         =		config_and_inputs
												_a	:    List[Any]         =		{'''input_ids''': input_ids, '''attention_mask''': input_mask}
												return config, inputs_dict



@require_torch
class 				UpperCAmelCase_		(			__lowercase				,       __lowercase				,       unittest.TestCase   ):
							"""simple docstring"""







							UpperCAmelCase__       :						Any	    =    False

							UpperCAmelCase__       :						Any	    =    (EsmForProteinFolding,) if is_torch_available() else ()
							UpperCAmelCase__       :						Union[str, Any]	    =    ()
							UpperCAmelCase__       :						int	    =    {} if is_torch_available() else {}
							UpperCAmelCase__       :						Optional[int]	    =    False


							def 				__lowercase				(			self       )    ->    List[Any]:
												_a	:    Optional[int]         =		EsmFoldModelTester(self       )
												_a	:    Dict         =		ConfigTester(self       ,							config_class=_a       ,							hidden_size=3_7       )


							def 				__lowercase				(			self       )    ->    List[str]:
												self.config_tester.run_common_tests()


							def 				__lowercase				(			self       )    ->    str:
												_a	:    str         =		self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_a       )


							@unittest.skip('''Does not support attention outputs'''       )
							def 				__lowercase				(			self       )    ->    int:
												pass


							@unittest.skip
							def 				__lowercase				(			self       )    ->    List[str]:
												pass


							@unittest.skip('''Esm does not support embedding resizing'''       )
							def 				__lowercase				(			self       )    ->    int:
												pass


							@unittest.skip('''Esm does not support embedding resizing'''       )
							def 				__lowercase				(			self       )    ->    Optional[Any]:
												pass


							@unittest.skip('''ESMFold does not support passing input embeds!'''       )
							def 				__lowercase				(			self       )    ->    int:
												pass


							@unittest.skip('''ESMFold does not support head pruning.'''       )
							def 				__lowercase				(			self       )    ->    str:
												pass


							@unittest.skip('''ESMFold does not support head pruning.'''       )
							def 				__lowercase				(			self       )    ->    Optional[Any]:
												pass


							@unittest.skip('''ESMFold does not support head pruning.'''       )
							def 				__lowercase				(			self       )    ->    Optional[Any]:
												pass


							@unittest.skip('''ESMFold does not support head pruning.'''       )
							def 				__lowercase				(			self       )    ->    Any:
												pass


							@unittest.skip('''ESMFold does not support head pruning.'''       )
							def 				__lowercase				(			self       )    ->    Tuple:
												pass


							@unittest.skip('''ESMFold does not output hidden states in the normal way.'''       )
							def 				__lowercase				(			self       )    ->    Tuple:
												pass


							@unittest.skip('''ESMfold does not output hidden states in the normal way.'''       )
							def 				__lowercase				(			self       )    ->    Tuple:
												pass


							@unittest.skip('''ESMFold only has one output format.'''       )
							def 				__lowercase				(			self       )    ->    Tuple:
												pass


							@unittest.skip('''This test doesn\'t work for ESMFold and doesn\'t test core functionality'''       )
							def 				__lowercase				(			self       )    ->    Dict:
												pass


							@unittest.skip('''ESMFold does not support input chunking.'''       )
							def 				__lowercase				(			self       )    ->    Tuple:
												pass


							@unittest.skip('''ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.'''       )
							def 				__lowercase				(			self       )    ->    Optional[Any]:
												pass


							@unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''       )
							def 				__lowercase				(			self       )    ->    List[str]:
												pass


							@unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''       )
							def 				__lowercase				(			self       )    ->    List[str]:
												pass


							@unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''       )
							def 				__lowercase				(			self       )    ->    List[Any]:
												pass


							@unittest.skip('''ESMFold doesn\'t support data parallel.'''       )
							def 				__lowercase				(			self       )    ->    Union[str, Any]:
												pass




							@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''       )
							def 				__lowercase				(			self       )    ->    Optional[Any]:
												pass



@require_torch
class 				UpperCAmelCase_		(			__lowercase   ):
							"""simple docstring"""







							@slow
							def 				__lowercase				(			self       )    ->    Optional[int]:
												_a	:    Dict         =		EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1'''       ).float()
												model.eval()
												_a	:    Tuple         =		torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]       )
												_a	:    Optional[Any]         =		model(_a       )['''positions''']
												_a	:    Union[str, Any]         =		torch.tensor([2.5828, 0.7993, -10.9334]       ,							dtype=torch.floataa       )
												self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]       ,							_a       ,							atol=1e-4       )       )



