






'''simple docstring'''

from __future__ import annotations

import unittest

from transformers import EsmConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
						import numpy
						import tensorflow as tf

						from transformers.models.esm.modeling_tf_esm import (
						    TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFEsmForMaskedLM,
						    TFEsmForSequenceClassification,
						    TFEsmForTokenClassification,
						    TFEsmModel,
						)






class 							__magic_name__     :







	"""simple docstring"""



	def __init__(							self  ,      _lowercase  ,      )						->  int:
							lowercase_      :   Dict        =			parent
							lowercase_      :   Dict        =			13
							lowercase_      :   Optional[Any]        =			7
							lowercase_      :   Union[str, Any]        =			True
							lowercase_      :   str        =			True
							lowercase_      :   Optional[Any]        =			True
							lowercase_      :   List[Any]        =			99
							lowercase_      :   Union[str, Any]        =			32
							lowercase_      :   str        =			2
							lowercase_      :   Optional[int]        =			4
							lowercase_      :   Optional[int]        =			37
							lowercase_      :   Any        =			'gelu'
							lowercase_      :   Tuple        =			0.1
							lowercase_      :   List[Any]        =			0.1
							lowercase_      :   int        =			512
							lowercase_      :   Tuple        =			16
							lowercase_      :   Any        =			2
							lowercase_      :   Tuple        =			0.02
							lowercase_      :   Union[str, Any]        =			3
							lowercase_      :   int        =			4
							lowercase_      :   Optional[Any]        =			None
	def        lowerCamelCase__		(							self       )						->  Any:
							lowercase_      :   Optional[int]        =			ids_tensor([self.batch_size, self.seq_length]  ,      self.vocab_size       )

							lowercase_      :   Tuple        =			None
							if self.use_input_mask:
													lowercase_      :   Tuple        =			random_attention_mask([self.batch_size, self.seq_length]       )

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

							lowercase_      :   str        =			EsmConfig(
							    vocab_size=self.vocab_size  ,      hidden_size=self.hidden_size  ,      num_hidden_layers=self.num_hidden_layers  ,      pad_token_id=1  ,      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  ,      )

							return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
	def        lowerCamelCase__		(							self       )						->  Union[str, Any]:
							(
							    lowercase_      
							)      :   Any        =			self.prepare_config_and_inputs()

							lowercase_      :   Dict        =			True
							lowercase_      :   Optional[Any]        =			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
							lowercase_      :   Union[str, Any]        =			ids_tensor([self.batch_size, self.seq_length]  ,      vocab_size=2       )

							return (
							    config,
							    input_ids,
							    input_mask,
							    sequence_labels,
							    token_labels,
							    choice_labels,
							    encoder_hidden_states,
							    encoder_attention_mask,
							)
	def        lowerCamelCase__		(							self  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase       )						->  str:
							lowercase_      :   int        =			TFEsmModel(config=_lowercase       )
							lowercase_      :   List[str]        =			{'input_ids': input_ids, 'attention_mask': input_mask}
							lowercase_      :   Tuple        =			model(_lowercase       )

							lowercase_      :   Dict        =			[input_ids, input_mask]
							lowercase_      :   List[Any]        =			model(_lowercase       )

							lowercase_      :   Tuple        =			model(_lowercase       )

							self.parent.assertEqual(result.last_hidden_state.shape  ,      (self.batch_size, self.seq_length, self.hidden_size)       )
	def        lowerCamelCase__		(							self  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      )						->  List[Any]:
							lowercase_      :   Union[str, Any]        =			True

							lowercase_      :   Tuple        =			TFEsmModel(config=_lowercase       )
							lowercase_      :   int        =			{
							    'input_ids': input_ids,
							    'attention_mask': input_mask,
							    'encoder_hidden_states': encoder_hidden_states,
							    'encoder_attention_mask': encoder_attention_mask,
							}
							lowercase_      :   Optional[int]        =			model(_lowercase       )

							lowercase_      :   str        =			[input_ids, input_mask]
							lowercase_      :   Union[str, Any]        =			model(_lowercase  ,      encoder_hidden_states=_lowercase       )

							# Also check the case where encoder outputs are not passed
							lowercase_      :   List[Any]        =			model(_lowercase  ,      attention_mask=_lowercase       )

							self.parent.assertEqual(result.last_hidden_state.shape  ,      (self.batch_size, self.seq_length, self.hidden_size)       )
	def        lowerCamelCase__		(							self  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase       )						->  Tuple:
							lowercase_      :   Optional[Any]        =			TFEsmForMaskedLM(config=_lowercase       )
							lowercase_      :   List[Any]        =			model([input_ids, input_mask]       )
							self.parent.assertEqual(result.logits.shape  ,      (self.batch_size, self.seq_length, self.vocab_size)       )
	def        lowerCamelCase__		(							self  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase       )						->  Union[str, Any]:
							lowercase_      :   Union[str, Any]        =			self.num_labels
							lowercase_      :   Dict        =			TFEsmForTokenClassification(config=_lowercase       )
							lowercase_      :   str        =			{'input_ids': input_ids, 'attention_mask': input_mask}
							lowercase_      :   int        =			model(_lowercase       )
							self.parent.assertEqual(result.logits.shape  ,      (self.batch_size, self.seq_length, self.num_labels)       )







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






@require_tf
class 							__magic_name__     (	UpperCAmelCase_,						UpperCAmelCase_,						unittest.TestCase				):







	"""simple docstring"""



	SCREAMING_SNAKE_CASE_					:      Optional[Any]        =  (
	    (
	        TFEsmModel,
	        TFEsmForMaskedLM,
	        TFEsmForSequenceClassification,
	        TFEsmForTokenClassification,
	    )
	    if is_tf_available()
	    else ()
	)
	SCREAMING_SNAKE_CASE_					:      Tuple        =  (
	    {
	        'feature-extraction': TFEsmModel,
	        'fill-mask': TFEsmForMaskedLM,
	        'text-classification': TFEsmForSequenceClassification,
	        'token-classification': TFEsmForTokenClassification,
	        'zero-shot': TFEsmForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	SCREAMING_SNAKE_CASE_					:      List[Any]        =  False
	SCREAMING_SNAKE_CASE_					:      Optional[Any]        =  False
	def        lowerCamelCase__		(							self       )						->  Any:
							lowercase_      :   Optional[int]        =			TFEsmModelTester(self       )
							lowercase_      :   Any        =			ConfigTester(self  ,      config_class=_lowercase  ,      hidden_size=37       )
	def        lowerCamelCase__		(							self       )						->  str:
							self.config_tester.run_common_tests()
	def        lowerCamelCase__		(							self       )						->  Any:
							lowercase_      :   int        =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_lowercase       )
	def        lowerCamelCase__		(							self       )						->  Tuple:
							lowercase_      :   Optional[Any]        =			self.model_tester.prepare_config_and_inputs_for_decoder()
							self.model_tester.create_and_check_model_as_decoder(*_lowercase       )
	def        lowerCamelCase__		(							self       )						->  Optional[int]:
							lowercase_      :   Dict        =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*_lowercase       )
	def        lowerCamelCase__		(							self       )						->  str:
							lowercase_      :   Tuple        =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*_lowercase       )
	@slow
	def        lowerCamelCase__		(							self       )						->  Optional[int]:
							for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													lowercase_      :   List[str]        =			TFEsmModel.from_pretrained(_lowercase       )
													self.assertIsNotNone(_lowercase       )
	@unittest.skip('Protein models do not support embedding resizing.'       )
	def        lowerCamelCase__		(							self       )						->  Optional[int]:
							pass
	@unittest.skip('Protein models do not support embedding resizing.'       )
	def        lowerCamelCase__		(							self       )						->  Tuple:
							pass







	def        lowerCamelCase__		(							self       )						->  List[str]:
							lowercase_      :   str        =			self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													lowercase_      :   List[Any]        =			model_class(_lowercase       )
													assert isinstance(model.get_input_embeddings()  ,      tf.keras.layers.Layer       )
													if model_class is TFEsmForMaskedLM:
																			# Output embedding test differs from the main test because they're a matrix, not a layer
																			lowercase_      :   Any        =			model.get_bias()
																			assert isinstance(_lowercase  ,      _lowercase       )
																			for k, v in name.items():
																									assert isinstance(_lowercase  ,      tf.Variable       )
													else:
																			lowercase_      :   Union[str, Any]        =			model.get_output_embeddings()
																			assert x is None
																			lowercase_      :   Tuple        =			model.get_bias()
																			assert name is None






@require_tf
class 							__magic_name__     (	unittest.TestCase				):







	"""simple docstring"""



	@slow
	def        lowerCamelCase__		(							self       )						->  Optional[int]:
							lowercase_      :   Optional[int]        =			TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'       )

							lowercase_      :   Any        =			tf.constant([[0, 1, 2, 3, 4, 5]]       )
							lowercase_      :   Tuple        =			model(_lowercase       )[0]
							lowercase_      :   int        =			[1, 6, 33]
							self.assertEqual(list(output.numpy().shape       )  ,      _lowercase       )
							# compare the actual values for a slice.
							lowercase_      :   Union[str, Any]        =			tf.constant(
							    [
							        [
							            [8.92_15_18, -10.58_98_14, -6.4_67_13_07],
							            [-6.3_96_71_56, -13.91_13_77, -1.1_21_19_15],
							            [-7.78_12_47, -13.95_15_57, -3.74_05_92],
							        ]
							    ]       )
							self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()  ,      expected_slice.numpy()  ,      atol=1E-2       )       )







	@slow
	def        lowerCamelCase__		(							self       )						->  Union[str, Any]:
							lowercase_      :   Dict        =			TFEsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'       )

							lowercase_      :   int        =			tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]       )
							lowercase_      :   str        =			model(_lowercase       )[0]
							# compare the actual values for a slice.
							lowercase_      :   Optional[int]        =			tf.constant(
							    [
							        [
							            [0.14_44_30_92, 0.54_12_53_27, 0.3_24_77_39],
							            [0.30_34_04_84, 0.00_52_66_76, 0.31_07_77_22],
							            [0.32_27_80_43, -0.24_98_70_96, 0.3_41_46_28],
							        ]
							    ]       )
							self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()  ,      expected_slice.numpy()  ,      atol=1E-4       )       )








'''simple docstring'''

import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible






class 							__magic_name__     (	unittest.TestCase				):







	"""simple docstring"""



	def        lowerCamelCase__		(							self       )						->  Optional[Any]:
							lowercase_      :   Tuple        =			[
							    'safety_checker/pytorch_model.bin',
							    'safety_checker/model.safetensors',
							    'vae/diffusion_pytorch_model.bin',
							    'vae/diffusion_pytorch_model.safetensors',
							    'text_encoder/pytorch_model.bin',
							    'text_encoder/model.safetensors',
							    'unet/diffusion_pytorch_model.bin',
							    'unet/diffusion_pytorch_model.safetensors',
							]
							self.assertTrue(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Dict:
							lowercase_      :   Tuple        =			[
							    'unet/diffusion_pytorch_model.bin',
							    'unet/diffusion_pytorch_model.safetensors',
							]
							self.assertTrue(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  str:
							lowercase_      :   int        =			[
							    'safety_checker/pytorch_model.bin',
							    'safety_checker/model.safetensors',
							    'vae/diffusion_pytorch_model.bin',
							    'vae/diffusion_pytorch_model.safetensors',
							    'text_encoder/pytorch_model.bin',
							    'text_encoder/model.safetensors',
							    'unet/diffusion_pytorch_model.bin',
							    # Removed: 'unet/diffusion_pytorch_model.safetensors',
							]
							self.assertFalse(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Any:
							lowercase_      :   Dict        =			[
							    'text_encoder/pytorch_model.bin',
							    'text_encoder/model.safetensors',
							]
							self.assertTrue(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Optional[Any]:
							lowercase_      :   List[Any]        =			[
							    'safety_checker/pytorch_model.bin',
							    'safety_checker/model.safetensors',
							    'vae/diffusion_pytorch_model.bin',
							    'vae/diffusion_pytorch_model.safetensors',
							    'text_encoder/pytorch_model.bin',
							    # Removed: 'text_encoder/model.safetensors',
							    'unet/diffusion_pytorch_model.bin',
							    'unet/diffusion_pytorch_model.safetensors',
							]
							self.assertFalse(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Optional[int]:
							lowercase_      :   str        =			[
							    'safety_checker/pytorch_model.fp16.bin',
							    'safety_checker/model.fp16.safetensors',
							    'vae/diffusion_pytorch_model.fp16.bin',
							    'vae/diffusion_pytorch_model.fp16.safetensors',
							    'text_encoder/pytorch_model.fp16.bin',
							    'text_encoder/model.fp16.safetensors',
							    'unet/diffusion_pytorch_model.fp16.bin',
							    'unet/diffusion_pytorch_model.fp16.safetensors',
							]
							lowercase_      :   Union[str, Any]        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  int:
							lowercase_      :   Optional[int]        =			[
							    'unet/diffusion_pytorch_model.fp16.bin',
							    'unet/diffusion_pytorch_model.fp16.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  int:
							# pass variant but use the non-variant filenames
							lowercase_      :   Optional[int]        =			[
							    'unet/diffusion_pytorch_model.bin',
							    'unet/diffusion_pytorch_model.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Union[str, Any]:
							lowercase_      :   int        =			[
							    'safety_checker/pytorch_model.fp16.bin',
							    'safety_checker/model.fp16.safetensors',
							    'vae/diffusion_pytorch_model.fp16.bin',
							    'vae/diffusion_pytorch_model.fp16.safetensors',
							    'text_encoder/pytorch_model.fp16.bin',
							    'text_encoder/model.fp16.safetensors',
							    'unet/diffusion_pytorch_model.fp16.bin',
							    # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertFalse(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  int:
							lowercase_      :   str        =			[
							    'text_encoder/pytorch_model.fp16.bin',
							    'text_encoder/model.fp16.safetensors',
							]
							lowercase_      :   str        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  List[str]:
							# pass variant but use the non-variant filenames
							lowercase_      :   List[Any]        =			[
							    'text_encoder/pytorch_model.bin',
							    'text_encoder/model.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )







	def        lowerCamelCase__		(							self       )						->  List[str]:
							lowercase_      :   Union[str, Any]        =			[
							    'safety_checker/pytorch_model.fp16.bin',
							    'safety_checker/model.fp16.safetensors',
							    'vae/diffusion_pytorch_model.fp16.bin',
							    'vae/diffusion_pytorch_model.fp16.safetensors',
							    'text_encoder/pytorch_model.fp16.bin',
							    # 'text_encoder/model.fp16.safetensors',
							    'unet/diffusion_pytorch_model.fp16.bin',
							    'unet/diffusion_pytorch_model.fp16.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertFalse(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
