



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   _A   :


						'''simple docstring'''





						def __init__(	self	:							Dict							,      lowerCamelCase	:							Any							,      ):






							'''simple docstring'''



							__lowercase		       =			parent
							__lowercase		       =			13
							__lowercase		       =			7
							__lowercase		       =			True
							__lowercase		       =			True
							__lowercase		       =			True
							__lowercase		       =			99
							__lowercase		       =			32
							__lowercase		       =			2
							__lowercase		       =			4
							__lowercase		       =			37
							__lowercase		       =			"gelu"
							__lowercase		       =			0.1
							__lowercase		       =			0.1
							__lowercase		       =			512
							__lowercase		       =			16
							__lowercase		       =			2
							__lowercase		       =			0.02
							__lowercase		       =			3
							__lowercase		       =			4
							__lowercase		       =			None



						def     _snake_case	(	self	:							int	):






							'''simple docstring'''



							__lowercase		       =			ids_tensor([self.batch_size, self.seq_length]							,      self.vocab_size	)

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

							__lowercase		       =			None
							__lowercase		       =			None
							__lowercase		       =			None
							if self.use_labels:
								__lowercase		       =			ids_tensor([self.batch_size]							,      self.type_sequence_label_size	)
								__lowercase		       =			ids_tensor([self.batch_size, self.seq_length]							,      self.num_labels	)
								__lowercase		       =			ids_tensor([self.batch_size]							,      self.num_choices	)

							__lowercase		       =			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     _snake_case	(	self	:							Tuple	):






							'''simple docstring'''



							(
							    (
							    __lowercase
							)       ,    (
							    __lowercase
							)       ,    (
							    __lowercase
							)       ,    (
							    __lowercase
							)       ,    (
							    __lowercase
							)       ,    (
							    __lowercase
							)       ,    
							)		       =			self.prepare_config_and_inputs()

							__lowercase		       =			True
							__lowercase		       =			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]	)
							__lowercase		       =			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     _snake_case	(	self	:							Union[str, Any]							,      lowerCamelCase	:							int							,      lowerCamelCase	:							Union[str, Any]							,      lowerCamelCase	:							Optional[int]							,      lowerCamelCase	:							Optional[Any]							,      lowerCamelCase	:							str							,      lowerCamelCase	:							str	):






							'''simple docstring'''



							__lowercase		       =			TFEsmModel(config=lowerCamelCase	)
							__lowercase		       =			{"input_ids": input_ids, "attention_mask": input_mask}
							__lowercase		       =			model(lowerCamelCase	)

							__lowercase		       =			[input_ids, input_mask]
							__lowercase		       =			model(lowerCamelCase	)

							__lowercase		       =			model(lowerCamelCase	)

							self.parent.assertEqual(result.last_hidden_state.shape							,      (self.batch_size, self.seq_length, self.hidden_size)	)



						def     _snake_case	(	self	:							Optional[Any]							,      lowerCamelCase	:							List[str]							,      lowerCamelCase	:							Tuple							,      lowerCamelCase	:							Optional[Any]							,      lowerCamelCase	:							Optional[int]							,      lowerCamelCase	:							List[str]							,      lowerCamelCase	:							Any							,      lowerCamelCase	:							int							,      lowerCamelCase	:							Any							,      ):






							'''simple docstring'''



							__lowercase		       =			True

							__lowercase		       =			TFEsmModel(config=lowerCamelCase	)
							__lowercase		       =			{
							    "input_ids": input_ids,
							    "attention_mask": input_mask,
							    "encoder_hidden_states": encoder_hidden_states,
							    "encoder_attention_mask": encoder_attention_mask,
							}
							__lowercase		       =			model(lowerCamelCase	)

							__lowercase		       =			[input_ids, input_mask]
							__lowercase		       =			model(lowerCamelCase							,      encoder_hidden_states=lowerCamelCase	)

							# Also check the case where encoder outputs are not passed
							__lowercase		       =			model(lowerCamelCase							,      attention_mask=lowerCamelCase	)

							self.parent.assertEqual(result.last_hidden_state.shape							,      (self.batch_size, self.seq_length, self.hidden_size)	)



						def     _snake_case	(	self	:							Dict							,      lowerCamelCase	:							Tuple							,      lowerCamelCase	:							Optional[Any]							,      lowerCamelCase	:							Dict							,      lowerCamelCase	:							Any							,      lowerCamelCase	:							Union[str, Any]							,      lowerCamelCase	:							List[str]	):






							'''simple docstring'''



							__lowercase		       =			TFEsmForMaskedLM(config=lowerCamelCase	)
							__lowercase		       =			model([input_ids, input_mask]	)
							self.parent.assertEqual(result.logits.shape							,      (self.batch_size, self.seq_length, self.vocab_size)	)



						def     _snake_case	(	self	:							Dict							,      lowerCamelCase	:							Dict							,      lowerCamelCase	:							str							,      lowerCamelCase	:							Union[str, Any]							,      lowerCamelCase	:							List[Any]							,      lowerCamelCase	:							Tuple							,      lowerCamelCase	:							int	):






							'''simple docstring'''



							__lowercase		       =			self.num_labels
							__lowercase		       =			TFEsmForTokenClassification(config=lowerCamelCase	)
							__lowercase		       =			{"input_ids": input_ids, "attention_mask": input_mask}
							__lowercase		       =			model(lowerCamelCase	)
							self.parent.assertEqual(result.logits.shape							,      (self.batch_size, self.seq_length, self.num_labels)	)



						def     _snake_case	(	self	:							Tuple	):






							'''simple docstring'''



							__lowercase		       =			self.prepare_config_and_inputs()
							(
							    (
							    __lowercase
							)       ,    (
							    __lowercase
							)       ,    (
							    __lowercase
							)       ,    (
							    __lowercase
							)       ,    (
							    __lowercase
							)       ,    (
							    __lowercase
							)       ,    
							)		       =			config_and_inputs
							__lowercase		       =			{"input_ids": input_ids, "attention_mask": input_mask}
							return config, inputs_dict

@require_tf
class   _A   ( _lowercase  ,						_lowercase  ,						unittest.TestCase	):


						'''simple docstring'''





						_snake_case			:     Dict      							=      (
						    (
						        TFEsmModel,
						        TFEsmForMaskedLM,
						        TFEsmForSequenceClassification,
						        TFEsmForTokenClassification,
						    )
						    if is_tf_available()
						    else ()
						)
						_snake_case			:     str      							=      (
						    {
						        """feature-extraction""": TFEsmModel,
						        """fill-mask""": TFEsmForMaskedLM,
						        """text-classification""": TFEsmForSequenceClassification,
						        """token-classification""": TFEsmForTokenClassification,
						        """zero-shot""": TFEsmForSequenceClassification,
						    }
						    if is_tf_available()
						    else {}
						)
						_snake_case			:     Union[str, Any]      							=      False
						_snake_case			:     List[Any]      							=      False



						def     _snake_case	(	self	:							Tuple	):






							'''simple docstring'''



							__lowercase		       =			TFEsmModelTester(self	)
							__lowercase		       =			ConfigTester(self							,      config_class=lowerCamelCase							,      hidden_size=37	)



						def     _snake_case	(	self	:							Any	):






							'''simple docstring'''



							self.config_tester.run_common_tests()



						def     _snake_case	(	self	:							str	):






							'''simple docstring'''



							__lowercase		       =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*lowerCamelCase	)



						def     _snake_case	(	self	:							Dict	):






							'''simple docstring'''



							__lowercase		       =			self.model_tester.prepare_config_and_inputs_for_decoder()
							self.model_tester.create_and_check_model_as_decoder(*lowerCamelCase	)



						def     _snake_case	(	self	:							str	):






							'''simple docstring'''



							__lowercase		       =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase	)



						def     _snake_case	(	self	:							Optional[int]	):






							'''simple docstring'''



							__lowercase		       =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*lowerCamelCase	)



						@slow
						def     _snake_case	(	self	:							int	):






							'''simple docstring'''



							for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								__lowercase		       =			TFEsmModel.from_pretrained(lowerCamelCase	)
								self.assertIsNotNone(lowerCamelCase	)



						@unittest.skip("Protein models do not support embedding resizing."	)
						def     _snake_case	(	self	:							int	):






							'''simple docstring'''



							pass



						@unittest.skip("Protein models do not support embedding resizing."	)
						def     _snake_case	(	self	:							Optional[Any]	):






							'''simple docstring'''



							pass



						def     _snake_case	(	self	:							Optional[int]	):






							'''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	)
								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		       =			model.get_bias()
									assert isinstance(lowerCamelCase							,      lowerCamelCase	)
									for k, v in name.items():
										assert isinstance(lowerCamelCase							,      tf.Variable	)
								else:
									__lowercase		       =			model.get_output_embeddings()
									assert x is None
									__lowercase		       =			model.get_bias()
									assert name is None

@require_tf
class   _A   ( unittest.TestCase	):


						'''simple docstring'''





						@slow
						def     _snake_case	(	self	:							List[str]	):






							'''simple docstring'''



							__lowercase		       =			TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"	)

							__lowercase		       =			tf.constant([[0, 1, 2, 3, 4, 5]]	)
							__lowercase		       =			model(lowerCamelCase	)[0]
							__lowercase		       =			[1, 6, 33]
							self.assertEqual(list(output.numpy().shape	)							,      lowerCamelCase	)
							# compare the actual values for a slice.
							__lowercase		       =			tf.constant(
							    [
							        [
							            [8.92_1518, -10.58_9814, -6.467_1307],
							            [-6.396_7156, -13.91_1377, -1.121_1915],
							            [-7.78_1247, -13.95_1557, -3.74_0592],
							        ]
							    ]	)
							self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()							,      expected_slice.numpy()							,      atol=1e-2	)	)



						@slow
						def     _snake_case	(	self	:							Optional[Any]	):






							'''simple docstring'''



							__lowercase		       =			TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"	)

							__lowercase		       =			tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]	)
							__lowercase		       =			model(lowerCamelCase	)[0]
							# compare the actual values for a slice.
							__lowercase		       =			tf.constant(
							    [
							        [
							            [0.1444_3092, 0.5412_5327, 0.324_7739],
							            [0.3034_0484, 0.0052_6676, 0.3107_7722],
							            [0.3227_8043, -0.2498_7096, 0.341_4628],
							        ]
							    ]	)
							self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()							,      expected_slice.numpy()							,      atol=1e-4	)	)








import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np

class   _A   :


						'''simple docstring'''





						def     _snake_case	(	self	:							Union[str, Any]	):






							'''simple docstring'''



							torch.manual_seed(0	)
							__lowercase		       =			TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"	)

							torch.manual_seed(0	)
							__lowercase		       =			AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"	)

							torch.manual_seed(0	)
							__lowercase		       =			UNetaDConditionModel(
							    sample_size=32							,      layers_per_block=1							,      block_out_channels=[32, 64]							,      down_block_types=[
							        "ResnetDownsampleBlock2D",
							        "SimpleCrossAttnDownBlock2D",
							    ]							,      mid_block_type="UNetMidBlock2DSimpleCrossAttn"							,      up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]							,      in_channels=3							,      out_channels=6							,      cross_attention_dim=32							,      encoder_hid_dim=32							,      attention_head_dim=8							,      addition_embed_type="text"							,      addition_embed_type_num_heads=2							,      cross_attention_norm="group_norm"							,      resnet_time_scale_shift="scale_shift"							,      act_fn="gelu"							,      )
							unet.set_attn_processor(AttnAddedKVProcessor()	)  # For reproducibility tests

							torch.manual_seed(0	)
							__lowercase		       =			DDPMScheduler(
							    num_train_timesteps=1_000							,      beta_schedule="squaredcos_cap_v2"							,      beta_start=0.0001							,      beta_end=0.02							,      thresholding=lowerCamelCase							,      dynamic_thresholding_ratio=0.95							,      sample_max_value=1.0							,      prediction_type="epsilon"							,      variance_type="learned_range"							,      )

							torch.manual_seed(0	)
							__lowercase		       =			IFWatermarker()

							return {
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "unet": unet,
							    "scheduler": scheduler,
							    "watermarker": watermarker,
							    "safety_checker": None,
							    "feature_extractor": None,
							}



						def     _snake_case	(	self	:							Tuple	):






							'''simple docstring'''



							torch.manual_seed(0	)
							__lowercase		       =			TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"	)

							torch.manual_seed(0	)
							__lowercase		       =			AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"	)

							torch.manual_seed(0	)
							__lowercase		       =			UNetaDConditionModel(
							    sample_size=32							,      layers_per_block=[1, 2]							,      block_out_channels=[32, 64]							,      down_block_types=[
							        "ResnetDownsampleBlock2D",
							        "SimpleCrossAttnDownBlock2D",
							    ]							,      mid_block_type="UNetMidBlock2DSimpleCrossAttn"							,      up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]							,      in_channels=6							,      out_channels=6							,      cross_attention_dim=32							,      encoder_hid_dim=32							,      attention_head_dim=8							,      addition_embed_type="text"							,      addition_embed_type_num_heads=2							,      cross_attention_norm="group_norm"							,      resnet_time_scale_shift="scale_shift"							,      act_fn="gelu"							,      class_embed_type="timestep"							,      mid_block_scale_factor=1.414							,      time_embedding_act_fn="gelu"							,      time_embedding_dim=32							,      )
							unet.set_attn_processor(AttnAddedKVProcessor()	)  # For reproducibility tests

							torch.manual_seed(0	)
							__lowercase		       =			DDPMScheduler(
							    num_train_timesteps=1_000							,      beta_schedule="squaredcos_cap_v2"							,      beta_start=0.0001							,      beta_end=0.02							,      thresholding=lowerCamelCase							,      dynamic_thresholding_ratio=0.95							,      sample_max_value=1.0							,      prediction_type="epsilon"							,      variance_type="learned_range"							,      )

							torch.manual_seed(0	)
							__lowercase		       =			DDPMScheduler(
							    num_train_timesteps=1_000							,      beta_schedule="squaredcos_cap_v2"							,      beta_start=0.0001							,      beta_end=0.02							,      )

							torch.manual_seed(0	)
							__lowercase		       =			IFWatermarker()

							return {
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "unet": unet,
							    "scheduler": scheduler,
							    "image_noising_scheduler": image_noising_scheduler,
							    "watermarker": watermarker,
							    "safety_checker": None,
							    "feature_extractor": None,
							}



						def     _snake_case	(	self	:							str	):






							'''simple docstring'''



							__lowercase		       =			self.get_dummy_components()
							__lowercase		       =			self.pipeline_class(**lowerCamelCase	)
							pipe.to(lowerCamelCase	)
							pipe.set_progress_bar_config(disable=lowerCamelCase	)

							__lowercase		       =			self.get_dummy_inputs(lowerCamelCase	)

							__lowercase		       =			inputs["prompt"]
							__lowercase		       =			inputs["generator"]
							__lowercase		       =			inputs["num_inference_steps"]
							__lowercase		       =			inputs["output_type"]

							if "image" in inputs:
								__lowercase		       =			inputs["image"]
							else:
								__lowercase		       =			None

							if "mask_image" in inputs:
								__lowercase		       =			inputs["mask_image"]
							else:
								__lowercase		       =			None

							if "original_image" in inputs:
								__lowercase		       =			inputs["original_image"]
							else:
								__lowercase		       =			None

							__lowercase       ,    __lowercase		       =			pipe.encode_prompt(lowerCamelCase	)

							# inputs with prompt converted to embeddings
							__lowercase		       =			{
							    "prompt_embeds": prompt_embeds,
							    "negative_prompt_embeds": negative_prompt_embeds,
							    "generator": generator,
							    "num_inference_steps": num_inference_steps,
							    "output_type": output_type,
							}

							if image is not None:
								__lowercase		       =			image

							if mask_image is not None:
								__lowercase		       =			mask_image

							if original_image is not None:
								__lowercase		       =			original_image

							# set all optional components to None
							for optional_component in pipe._optional_components:
								setattr(lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase	)

							__lowercase		       =			pipe(**lowerCamelCase	)[0]

							with tempfile.TemporaryDirectory() as tmpdir:
								pipe.save_pretrained(lowerCamelCase	)
								__lowercase		       =			self.pipeline_class.from_pretrained(lowerCamelCase	)
								pipe_loaded.to(lowerCamelCase	)
								pipe_loaded.set_progress_bar_config(disable=lowerCamelCase	)

							pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()	)  # For reproducibility tests

							for optional_component in pipe._optional_components:
								self.assertTrue(
								    getattr(lowerCamelCase							,      lowerCamelCase	) is None							,      f"""`{optional_component}` did not stay set to None after loading."""							,      )

							__lowercase		       =			self.get_dummy_inputs(lowerCamelCase	)

							__lowercase		       =			inputs["generator"]
							__lowercase		       =			inputs["num_inference_steps"]
							__lowercase		       =			inputs["output_type"]

							# inputs with prompt converted to embeddings
							__lowercase		       =			{
							    "prompt_embeds": prompt_embeds,
							    "negative_prompt_embeds": negative_prompt_embeds,
							    "generator": generator,
							    "num_inference_steps": num_inference_steps,
							    "output_type": output_type,
							}

							if image is not None:
								__lowercase		       =			image

							if mask_image is not None:
								__lowercase		       =			mask_image

							if original_image is not None:
								__lowercase		       =			original_image

							__lowercase		       =			pipe_loaded(**lowerCamelCase	)[0]

							__lowercase		       =			np.abs(to_np(lowerCamelCase	) - to_np(lowerCamelCase	)	).max()
							self.assertLess(lowerCamelCase							,      1e-4	)



						def     _snake_case	(	self	:							Optional[Any]	):






							'''simple docstring'''



							__lowercase		       =			self.get_dummy_components()
							__lowercase		       =			self.pipeline_class(**lowerCamelCase	)
							pipe.to(lowerCamelCase	)
							pipe.set_progress_bar_config(disable=lowerCamelCase	)

							__lowercase		       =			self.get_dummy_inputs(lowerCamelCase	)
							__lowercase		       =			pipe(**lowerCamelCase	)[0]

							with tempfile.TemporaryDirectory() as tmpdir:
								pipe.save_pretrained(lowerCamelCase	)
								__lowercase		       =			self.pipeline_class.from_pretrained(lowerCamelCase	)
								pipe_loaded.to(lowerCamelCase	)
								pipe_loaded.set_progress_bar_config(disable=lowerCamelCase	)

							pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()	)  # For reproducibility tests

							__lowercase		       =			self.get_dummy_inputs(lowerCamelCase	)
							__lowercase		       =			pipe_loaded(**lowerCamelCase	)[0]

							__lowercase		       =			np.abs(to_np(lowerCamelCase	) - to_np(lowerCamelCase	)	).max()
							self.assertLess(lowerCamelCase							,      1e-4	)



