
from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
					import tensorflow as tf

					from transformers import (
					    TFDeiTForImageClassification,
					    TFDeiTForImageClassificationWithTeacher,
					    TFDeiTForMaskedImageModeling,
					    TFDeiTModel,
					)
					from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					from PIL import Image

					from transformers import DeiTImageProcessor







class  UpperCamelCase__		:







	'''simple docstring'''





	def __init__(			self				,       UpperCamelCase__				,       UpperCamelCase__=13				,       UpperCamelCase__=30				,       UpperCamelCase__=2				,       UpperCamelCase__=3				,       UpperCamelCase__=True				,       UpperCamelCase__=True				,       UpperCamelCase__=32				,       UpperCamelCase__=2				,       UpperCamelCase__=4				,       UpperCamelCase__=37				,       UpperCamelCase__="gelu"				,       UpperCamelCase__=0.1				,       UpperCamelCase__=0.1				,       UpperCamelCase__=10				,       UpperCamelCase__=0.0_2				,       UpperCamelCase__=3				,       UpperCamelCase__=None				,       UpperCamelCase__=2				,       ):
					A__  :	Any         =							parent
					A__  :	str         =							batch_size
					A__  :	Optional[int]         =							image_size
					A__  :	Tuple         =							patch_size
					A__  :	Dict         =							num_channels
					A__  :	Optional[Any]         =							is_training
					A__  :	List[str]         =							use_labels
					A__  :	int         =							hidden_size
					A__  :	Optional[Any]         =							num_hidden_layers
					A__  :	Union[str, Any]         =							num_attention_heads
					A__  :	str         =							intermediate_size
					A__  :	Any         =							hidden_act
					A__  :	List[str]         =							hidden_dropout_prob
					A__  :	int         =							attention_probs_dropout_prob
					A__  :	List[str]         =							type_sequence_label_size
					A__  :	Dict         =							initializer_range
					A__  :	Any         =							scope
					A__  :	Dict         =							encoder_stride

					# in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
					A__  :	Dict         =							(image_size // patch_size) ** 2
					A__  :	Dict         =							num_patches + 2





	def   __snake_case				(			self		):
					A__  :	int         =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

					A__  :	Optional[Any]         =							None
					if self.use_labels:
									A__  :	Tuple         =							ids_tensor([self.batch_size]				,       self.type_sequence_label_size		)

					A__  :	List[str]         =							self.get_config()

					return config, pixel_values, labels





	def   __snake_case				(			self		):
					return DeiTConfig(
					    image_size=self.image_size				,       patch_size=self.patch_size				,       num_channels=self.num_channels				,       hidden_size=self.hidden_size				,       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				,       is_decoder=UpperCamelCase__				,       initializer_range=self.initializer_range				,       encoder_stride=self.encoder_stride				,       )





	def   __snake_case				(			self				,       UpperCamelCase__				,       UpperCamelCase__				,       UpperCamelCase__		):
					A__  :	Dict         =							TFDeiTModel(config=UpperCamelCase__		)
					A__  :	int         =							model(UpperCamelCase__		)
					self.parent.assertEqual(result.last_hidden_state.shape				,       (self.batch_size, self.seq_length, self.hidden_size)		)





	def   __snake_case				(			self				,       UpperCamelCase__				,       UpperCamelCase__				,       UpperCamelCase__		):
					A__  :	List[Any]         =							TFDeiTForMaskedImageModeling(config=UpperCamelCase__		)
					A__  :	Tuple         =							model(UpperCamelCase__		)
					self.parent.assertEqual(
					    result.reconstruction.shape				,       (self.batch_size, self.num_channels, self.image_size, self.image_size)		)

					# test greyscale images
					A__  :	Optional[Any]         =							1
					A__  :	Any         =							TFDeiTForMaskedImageModeling(UpperCamelCase__		)

					A__  :	Dict         =							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
					A__  :	List[Any]         =							model(UpperCamelCase__		)
					self.parent.assertEqual(result.reconstruction.shape				,       (self.batch_size, 1, self.image_size, self.image_size)		)





	def   __snake_case				(			self				,       UpperCamelCase__				,       UpperCamelCase__				,       UpperCamelCase__		):
					A__  :	List[Any]         =							self.type_sequence_label_size
					A__  :	str         =							TFDeiTForImageClassification(UpperCamelCase__		)
					A__  :	Any         =							model(UpperCamelCase__				,       labels=UpperCamelCase__		)
					self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.type_sequence_label_size)		)

					# test greyscale images
					A__  :	Optional[Any]         =							1
					A__  :	Optional[int]         =							TFDeiTForImageClassification(UpperCamelCase__		)

					A__  :	Any         =							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
					A__  :	Optional[int]         =							model(UpperCamelCase__				,       labels=UpperCamelCase__		)
					self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.type_sequence_label_size)		)





	def   __snake_case				(			self		):
					A__  :	Dict         =							self.prepare_config_and_inputs()
					A__       ,							A__       ,							A__  :	Tuple         =							config_and_inputs
					A__  :	Optional[Any]         =							{'''pixel_values''': pixel_values}
					return config, inputs_dict







@require_tf
class  UpperCamelCase__		(    SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			unittest.TestCase   ):







	'''simple docstring'''
	_lowerCAmelCase							       =				(
	    (
	        TFDeiTModel,
	        TFDeiTForImageClassification,
	        TFDeiTForImageClassificationWithTeacher,
	        TFDeiTForMaskedImageModeling,
	    )
	    if is_tf_available()
	    else ()
	)
	_lowerCAmelCase							       =				(
	    {
	        "feature-extraction": TFDeiTModel,
	        "image-classification": (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
	    }
	    if is_tf_available()
	    else {}
	)

	_lowerCAmelCase							       =				False
	_lowerCAmelCase							       =				False
	_lowerCAmelCase							       =				False
	_lowerCAmelCase							       =				False





	def   __snake_case				(			self		):
					A__  :	int         =							TFDeiTModelTester(self		)
					A__  :	List[Any]         =							ConfigTester(self				,       config_class=UpperCamelCase__				,       has_text_modality=UpperCamelCase__				,       hidden_size=37		)





	def   __snake_case				(			self		):
					self.config_tester.run_common_tests()





	@unittest.skip(reason='''DeiT does not use inputs_embeds'''		)
	def   __snake_case				(			self		):
					pass





	def   __snake_case				(			self		):
					A__       ,							A__  :	str         =							self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
									A__  :	List[Any]         =							model_class(UpperCamelCase__		)
									self.assertIsInstance(model.get_input_embeddings()				,       (tf.keras.layers.Layer)		)
									A__  :	List[Any]         =							model.get_output_embeddings()
									self.assertTrue(x is None or isinstance(UpperCamelCase__				,       tf.keras.layers.Dense		)		)





	def   __snake_case				(			self		):
					A__       ,							A__  :	Dict         =							self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
									A__  :	List[Any]         =							model_class(UpperCamelCase__		)
									A__  :	Optional[Any]         =							inspect.signature(model.call		)
									# signature.parameters is an OrderedDict => so arg_names order is deterministic
									A__  :	List[str]         =							[*signature.parameters.keys()]

									A__  :	Dict         =							['''pixel_values''']
									self.assertListEqual(arg_names[:1]				,       UpperCamelCase__		)





	def   __snake_case				(			self		):
					A__  :	Optional[Any]         =							self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*UpperCamelCase__		)





	def   __snake_case				(			self		):
					A__  :	Optional[int]         =							self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_image_modeling(*UpperCamelCase__		)





	def   __snake_case				(			self		):
					A__  :	str         =							self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_image_classification(*UpperCamelCase__		)





	def   __snake_case				(			self				,       UpperCamelCase__				,       UpperCamelCase__				,       UpperCamelCase__=False		):
					A__  :	Any         =							super()._prepare_for_class(UpperCamelCase__				,       UpperCamelCase__				,       return_labels=UpperCamelCase__		)

					if return_labels:
									if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call		).parameters:
													del inputs_dict["labels"]

					return inputs_dict





	@slow
	def   __snake_case				(			self		):
					for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									A__  :	List[str]         =							TFDeiTModel.from_pretrained(UpperCamelCase__		)
									self.assertIsNotNone(UpperCamelCase__		)




def 				SCREAMING_SNAKE_CASE     (      ) ->       str:

				"""simple docstring"""







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







@require_tf
@require_vision
class  UpperCamelCase__		(    unittest.TestCase   ):







	'''simple docstring'''





	@cached_property
	def   __snake_case				(			self		):
					return (
					    DeiTImageProcessor.from_pretrained('''facebook/deit-base-distilled-patch16-224'''		)
					    if is_vision_available()
					    else None
					)





	@slow
	def   __snake_case				(			self		):
					A__  :	Any         =							TFDeiTForImageClassificationWithTeacher.from_pretrained('''facebook/deit-base-distilled-patch16-224'''		)

					A__  :	Optional[int]         =							self.default_image_processor
					A__  :	Dict         =							prepare_img()
					A__  :	Any         =							image_processor(images=UpperCamelCase__				,       return_tensors='''tf'''		)

					# forward pass
					A__  :	str         =							model(**UpperCamelCase__		)

					# verify the logits
					A__  :	Dict         =							tf.TensorShape((1, 1000)		)
					self.assertEqual(outputs.logits.shape				,       UpperCamelCase__		)

					A__  :	List[str]         =							tf.constant([-1.0_2_6_6, 0.1_9_1_2, -1.2_8_6_1]		)

					self.assertTrue(np.allclose(outputs.logits[0, :3]				,       UpperCamelCase__				,       atol=1e-4		)		)

from dataclasses import asdict, dataclass
from typing import Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_SCREAMING_SNAKE_CASE    :      Dict          =			logging.get_logger(__name__)

# TODO Update this
_SCREAMING_SNAKE_CASE    :      Optional[int]          =			{
    'facebook/esm-1b': 'https://huggingface.co/facebook/esm-1b/resolve/main/config.json',
    # See all ESM models at https://huggingface.co/models?filter=esm
}







class  UpperCamelCase__		(    SCREAMING_SNAKE_CASE_   ):







	'''simple docstring'''
	_lowerCAmelCase							       =				"esm"





	def __init__(			self				,       UpperCamelCase__=None				,       UpperCamelCase__=None				,       UpperCamelCase__=None				,       UpperCamelCase__=768				,       UpperCamelCase__=12				,       UpperCamelCase__=12				,       UpperCamelCase__=3072				,       UpperCamelCase__=0.1				,       UpperCamelCase__=0.1				,       UpperCamelCase__=1026				,       UpperCamelCase__=0.0_2				,       UpperCamelCase__=1e-12				,       UpperCamelCase__="absolute"				,       UpperCamelCase__=True				,       UpperCamelCase__=None				,       UpperCamelCase__=False				,       UpperCamelCase__=False				,       UpperCamelCase__=None				,       UpperCamelCase__=None				,       **UpperCamelCase__				,       ):
					super().__init__(pad_token_id=UpperCamelCase__				,       mask_token_id=UpperCamelCase__				,       **UpperCamelCase__		)

					A__  :	Optional[Any]         =							vocab_size
					A__  :	int         =							hidden_size
					A__  :	List[str]         =							num_hidden_layers
					A__  :	Tuple         =							num_attention_heads
					A__  :	str         =							intermediate_size
					A__  :	List[str]         =							hidden_dropout_prob
					A__  :	Optional[Any]         =							attention_probs_dropout_prob
					A__  :	int         =							max_position_embeddings
					A__  :	List[str]         =							initializer_range
					A__  :	List[Any]         =							layer_norm_eps
					A__  :	int         =							position_embedding_type
					A__  :	Optional[Any]         =							use_cache
					A__  :	Optional[int]         =							emb_layer_norm_before
					A__  :	List[str]         =							token_dropout
					A__  :	Tuple         =							is_folding_model
					if is_folding_model:
									if esmfold_config is None:
													logger.info('''No esmfold_config supplied for folding model, using default values.'''		)
													A__  :	List[Any]         =							EsmFoldConfig()
									elif isinstance(UpperCamelCase__				,       UpperCamelCase__		):
													A__  :	Optional[int]         =							EsmFoldConfig(**UpperCamelCase__		)
									A__  :	int         =							esmfold_config
									if vocab_list is None:
													logger.warning('''No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!'''		)
													A__  :	Any         =							get_default_vocab_list()
									else:
													A__  :	Dict         =							vocab_list
					else:
									A__  :	Optional[Any]         =							None
									A__  :	Tuple         =							None
					if self.esmfold_config is not None and getattr(self.esmfold_config				,       '''use_esm_attn_map'''				,       UpperCamelCase__		):
									raise ValueError('''The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!'''		)





	def   __snake_case				(			self		):
					A__  :	Optional[int]         =							super().to_dict()
					if isinstance(self.esmfold_config				,       UpperCamelCase__		):
									A__  :	Dict         =							self.esmfold_config.to_dict()
					return output







@dataclass
class  UpperCamelCase__		:







	'''simple docstring'''
	_lowerCAmelCase							       =				None
	_lowerCAmelCase							       =				True
	_lowerCAmelCase							       =				False
	_lowerCAmelCase							       =				False
	_lowerCAmelCase							       =				False
	_lowerCAmelCase							       =				0

	_lowerCAmelCase							       =				True
	_lowerCAmelCase							       =				False

	_lowerCAmelCase							       =				128
	_lowerCAmelCase							       =				None





	def   __snake_case				(			self		):
					if self.trunk is None:
									A__  :	Tuple         =							TrunkConfig()
					elif isinstance(self.trunk				,       UpperCamelCase__		):
									A__  :	List[Any]         =							TrunkConfig(**self.trunk		)





	def   __snake_case				(			self		):
					A__  :	Optional[int]         =							asdict(self		)
					A__  :	int         =							self.trunk.to_dict()
					return output







@dataclass
class  UpperCamelCase__		:







	'''simple docstring'''
	_lowerCAmelCase							       =				48
	_lowerCAmelCase							       =				1_024
	_lowerCAmelCase							       =				128
	_lowerCAmelCase							       =				32
	_lowerCAmelCase							       =				32
	_lowerCAmelCase							       =				32
	_lowerCAmelCase							       =				0
	_lowerCAmelCase							       =				0
	_lowerCAmelCase							       =				False
	_lowerCAmelCase							       =				4
	_lowerCAmelCase							       =				128
	_lowerCAmelCase							       =				None





	def   __snake_case				(			self		):
					if self.structure_module is None:
									A__  :	str         =							StructureModuleConfig()
					elif isinstance(self.structure_module				,       UpperCamelCase__		):
									A__  :	str         =							StructureModuleConfig(**self.structure_module		)

					if self.max_recycles <= 0:
									raise ValueError(F"`max_recycles` should be positive, got {self.max_recycles}."		)
					if self.sequence_state_dim % self.sequence_state_dim != 0:
									raise ValueError(
									    '''`sequence_state_dim` should be a round multiple of `sequence_state_dim`, got'''
									    F" {self.sequence_state_dim} and {self.sequence_state_dim}."		)
					if self.pairwise_state_dim % self.pairwise_state_dim != 0:
									raise ValueError(
									    '''`pairwise_state_dim` should be a round multiple of `pairwise_state_dim`, got'''
									    F" {self.pairwise_state_dim} and {self.pairwise_state_dim}."		)

					A__  :	Tuple         =							self.sequence_state_dim // self.sequence_head_width
					A__  :	int         =							self.pairwise_state_dim // self.pairwise_head_width

					if self.sequence_state_dim != sequence_num_heads * self.sequence_head_width:
									raise ValueError(
									    '''`sequence_state_dim` should be equal to `sequence_num_heads * sequence_head_width, got'''
									    F" {self.sequence_state_dim} != {sequence_num_heads} * {self.sequence_head_width}."		)
					if self.pairwise_state_dim != pairwise_num_heads * self.pairwise_head_width:
									raise ValueError(
									    '''`pairwise_state_dim` should be equal to `pairwise_num_heads * pairwise_head_width, got'''
									    F" {self.pairwise_state_dim} != {pairwise_num_heads} * {self.pairwise_head_width}."		)
					if self.pairwise_state_dim % 2 != 0:
									raise ValueError(F"`pairwise_state_dim` should be even, got {self.pairwise_state_dim}."		)

					if self.dropout >= 0.4:
									raise ValueError(F"`dropout` should not be greater than 0.4, got {self.dropout}."		)





	def   __snake_case				(			self		):
					A__  :	List[Any]         =							asdict(self		)
					A__  :	Optional[int]         =							self.structure_module.to_dict()
					return output







@dataclass
class  UpperCamelCase__		:







	'''simple docstring'''
	_lowerCAmelCase							       =				384
	_lowerCAmelCase							       =				128
	_lowerCAmelCase							       =				16
	_lowerCAmelCase							       =				128
	_lowerCAmelCase							       =				12
	_lowerCAmelCase							       =				4
	_lowerCAmelCase							       =				8
	_lowerCAmelCase							       =				0.1
	_lowerCAmelCase							       =				8
	_lowerCAmelCase							       =				1
	_lowerCAmelCase							       =				2
	_lowerCAmelCase							       =				7
	_lowerCAmelCase							       =				10
	_lowerCAmelCase							       =				1e-8
	_lowerCAmelCase							       =				1e5





	def   __snake_case				(			self		):
					return asdict(self		)




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

				"""simple docstring"""







				return (
				    "<cls>",
				    "<pad>",
				    "<eos>",
				    "<unk>",
				    "L",
				    "A",
				    "G",
				    "V",
				    "S",
				    "E",
				    "R",
				    "T",
				    "I",
				    "D",
				    "P",
				    "K",
				    "Q",
				    "N",
				    "F",
				    "Y",
				    "M",
				    "H",
				    "W",
				    "C",
				    "X",
				    "B",
				    "U",
				    "Z",
				    "O",
				    ".",
				    "-",
				    "<null_1>",
				    "<mask>",
				)