



'''simple docstring'''
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 import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
					from transformers.models.esm.modeling_esm import (
					    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
					    EsmEmbeddings,
					    create_position_ids_from_input_ids,
					)






class 							lowercase__	:





	def __init__( self					:				str ,lowerCamelCase__					:				Tuple ,lowerCamelCase__					:				str=13 ,lowerCamelCase__					:				Any=7 ,lowerCamelCase__					:				int=False ,lowerCamelCase__					:				Optional[Any]=True ,lowerCamelCase__					:				int=False ,lowerCamelCase__					:				str=True ,lowerCamelCase__					:				List[Any]=33 ,lowerCamelCase__					:				List[Any]=32 ,lowerCamelCase__					:				str=5 ,lowerCamelCase__					:				str=4 ,lowerCamelCase__					:				List[Any]=37 ,lowerCamelCase__					:				Union[str, Any]="gelu" ,lowerCamelCase__					:				List[str]=0.1 ,lowerCamelCase__					:				Tuple=0.1 ,lowerCamelCase__					:				str=512 ,lowerCamelCase__					:				Optional[int]=16 ,lowerCamelCase__					:				Optional[Any]=2 ,lowerCamelCase__					:				str=0.0_2 ,lowerCamelCase__					:				Any=3 ,lowerCamelCase__					:				Union[str, Any]=4 ,lowerCamelCase__					:				Tuple=None ,):







							'''simple docstring'''
							_UpperCamelCase					:		str 				=   parent
							_UpperCamelCase					:		int 				=   batch_size
							_UpperCamelCase					:		Dict 				=   seq_length
							_UpperCamelCase					:		Dict 				=   is_training
							_UpperCamelCase					:		Optional[int] 				=   use_input_mask
							_UpperCamelCase					:		Optional[int] 				=   use_token_type_ids
							_UpperCamelCase					:		Tuple 				=   use_labels
							_UpperCamelCase					:		Any 				=   vocab_size
							_UpperCamelCase					:		Dict 				=   hidden_size
							_UpperCamelCase					:		Union[str, Any] 				=   num_hidden_layers
							_UpperCamelCase					:		int 				=   num_attention_heads
							_UpperCamelCase					:		Dict 				=   intermediate_size
							_UpperCamelCase					:		str 				=   hidden_act
							_UpperCamelCase					:		int 				=   hidden_dropout_prob
							_UpperCamelCase					:		List[Any] 				=   attention_probs_dropout_prob
							_UpperCamelCase					:		List[Any] 				=   max_position_embeddings
							_UpperCamelCase					:		Tuple 				=   type_vocab_size
							_UpperCamelCase					:		List[str] 				=   type_sequence_label_size
							_UpperCamelCase					:		str 				=   initializer_range
							_UpperCamelCase					:		Union[str, Any] 				=   num_labels
							_UpperCamelCase					:		Union[str, Any] 				=   num_choices
							_UpperCamelCase					:		Any 				=   scope





	def 	UpperCamelCase_  ( self					:				Dict ):







							'''simple docstring'''
							_UpperCamelCase					:		int 				=   ids_tensor([self.batch_size, self.seq_length] ,self.vocab_size )

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

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

							_UpperCamelCase					:		Union[str, Any] 				=   self.get_config()

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





	def 	UpperCamelCase_  ( self					:				List[str] ):







							'''simple docstring'''
							return EsmConfig(
							    vocab_size=self.vocab_size ,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 ,)





	def 	UpperCamelCase_  ( self					:				Any ,lowerCamelCase__					:				Tuple ,lowerCamelCase__					:				Dict ,lowerCamelCase__					:				int ,lowerCamelCase__					:				int ,lowerCamelCase__					:				Tuple ,lowerCamelCase__					:				List[Any] ):







							'''simple docstring'''
							_UpperCamelCase					:		Union[str, Any] 				=   EsmModel(config=lowerCamelCase__ )
							model.to(lowerCamelCase__ )
							model.eval()
							_UpperCamelCase					:		Union[str, Any] 				=   model(lowerCamelCase__ ,attention_mask=lowerCamelCase__ )
							_UpperCamelCase					:		Any 				=   model(lowerCamelCase__ )
							_UpperCamelCase					:		Union[str, Any] 				=   model(lowerCamelCase__ )

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





	def 	UpperCamelCase_  ( self					:				List[Any] ,lowerCamelCase__					:				str ,lowerCamelCase__					:				Optional[Any] ,lowerCamelCase__					:				List[str] ,lowerCamelCase__					:				List[str] ,lowerCamelCase__					:				Optional[int] ,lowerCamelCase__					:				str ):







							'''simple docstring'''
							_UpperCamelCase					:		List[str] 				=   EsmForMaskedLM(config=lowerCamelCase__ )
							model.to(lowerCamelCase__ )
							model.eval()
							_UpperCamelCase					:		Tuple 				=   model(lowerCamelCase__ ,attention_mask=lowerCamelCase__ ,labels=lowerCamelCase__ )
							self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.seq_length, self.vocab_size) )





	def 	UpperCamelCase_  ( self					:				int ,lowerCamelCase__					:				List[Any] ,lowerCamelCase__					:				Dict ,lowerCamelCase__					:				Union[str, Any] ,lowerCamelCase__					:				Dict ,lowerCamelCase__					:				List[Any] ,lowerCamelCase__					:				Optional[int] ):







							'''simple docstring'''
							_UpperCamelCase					:		Any 				=   self.num_labels
							_UpperCamelCase					:		Optional[int] 				=   EsmForTokenClassification(config=lowerCamelCase__ )
							model.to(lowerCamelCase__ )
							model.eval()
							_UpperCamelCase					:		Union[str, Any] 				=   model(lowerCamelCase__ ,attention_mask=lowerCamelCase__ ,labels=lowerCamelCase__ )
							self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.seq_length, self.num_labels) )





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







							'''simple docstring'''
							_UpperCamelCase					:		str 				=   self.prepare_config_and_inputs()
							(
							    (
							    _UpperCamelCase
							)				,			(
							    _UpperCamelCase
							)				,			(
							    _UpperCamelCase
							)				,			(
							    _UpperCamelCase
							)				,			(
							    _UpperCamelCase
							)				,			(
							    _UpperCamelCase
							)				,			
							)					:		Optional[Any] 				=   config_and_inputs
							_UpperCamelCase					:		int 				=   {'input_ids': input_ids, 'attention_mask': input_mask}
							return config, inputs_dict






@require_torch
class 							lowercase__	(							lowercase						,					lowercase						,					unittest.TestCase   ):
	lowercase__  				=			False

	lowercase__  				=			(
	    (
	        EsmForMaskedLM,
	        EsmModel,
	        EsmForSequenceClassification,
	        EsmForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	lowercase__  				=			()
	lowercase__  				=			(
	    {
	        """feature-extraction""": EsmModel,
	        """fill-mask""": EsmForMaskedLM,
	        """text-classification""": EsmForSequenceClassification,
	        """token-classification""": EsmForTokenClassification,
	        """zero-shot""": EsmForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	lowercase__  				=			True





	def 	UpperCamelCase_  ( self					:				Dict ):







							'''simple docstring'''
							_UpperCamelCase					:		Dict 				=   EsmModelTester(self )
							_UpperCamelCase					:		List[str] 				=   ConfigTester(self ,config_class=lowerCamelCase__ ,hidden_size=37 )





	def 	UpperCamelCase_  ( self					:				int ):







							'''simple docstring'''
							self.config_tester.run_common_tests()





	def 	UpperCamelCase_  ( self					:				Dict ):







							'''simple docstring'''
							_UpperCamelCase					:		Tuple 				=   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*lowerCamelCase__ )





	def 	UpperCamelCase_  ( self					:				Optional[int] ):







							'''simple docstring'''
							_UpperCamelCase					:		int 				=   self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
													_UpperCamelCase					:		int 				=   type
													self.model_tester.create_and_check_model(*lowerCamelCase__ )





	def 	UpperCamelCase_  ( self					:				Dict ):







							'''simple docstring'''
							_UpperCamelCase					:		List[str] 				=   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase__ )





	def 	UpperCamelCase_  ( self					:				int ):







							'''simple docstring'''
							_UpperCamelCase					:		str 				=   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*lowerCamelCase__ )





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







							'''simple docstring'''
							for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													_UpperCamelCase					:		Optional[int] 				=   EsmModel.from_pretrained(lowerCamelCase__ )
													self.assertIsNotNone(lowerCamelCase__ )





	def 	UpperCamelCase_  ( self					:				Optional[Any] ):







							'''simple docstring'''
							_UpperCamelCase					:		Optional[Any] 				=   self.model_tester.prepare_config_and_inputs()[0]
							_UpperCamelCase					:		Optional[int] 				=   EsmEmbeddings(config=lowerCamelCase__ )

							_UpperCamelCase					:		List[str] 				=   torch.as_tensor([[12, 31, 13, model.padding_idx]] )
							_UpperCamelCase					:		Optional[Any] 				=   torch.as_tensor(
							    [
							        [
							            0 + model.padding_idx + 1,
							            1 + model.padding_idx + 1,
							            2 + model.padding_idx + 1,
							            model.padding_idx,
							        ]
							    ] )
							_UpperCamelCase					:		int 				=   create_position_ids_from_input_ids(lowerCamelCase__ ,model.padding_idx )
							self.assertEqual(position_ids.shape ,expected_positions.shape )
							self.assertTrue(torch.all(torch.eq(lowerCamelCase__ ,lowerCamelCase__ ) ) )





	def 	UpperCamelCase_  ( self					:				Any ):







							'''simple docstring'''
							_UpperCamelCase					:		str 				=   self.model_tester.prepare_config_and_inputs()[0]
							_UpperCamelCase					:		List[str] 				=   EsmEmbeddings(config=lowerCamelCase__ )

							_UpperCamelCase					:		List[Any] 				=   torch.empty(2 ,4 ,30 )
							_UpperCamelCase					:		Optional[int] 				=   [
							    0 + embeddings.padding_idx + 1,
							    1 + embeddings.padding_idx + 1,
							    2 + embeddings.padding_idx + 1,
							    3 + embeddings.padding_idx + 1,
							]
							_UpperCamelCase					:		Union[str, Any] 				=   torch.as_tensor([expected_single_positions, expected_single_positions] )
							_UpperCamelCase					:		Dict 				=   embeddings.create_position_ids_from_inputs_embeds(lowerCamelCase__ )
							self.assertEqual(position_ids.shape ,expected_positions.shape )
							self.assertTrue(torch.all(torch.eq(lowerCamelCase__ ,lowerCamelCase__ ) ) )





	@unittest.skip('Esm does not support embedding resizing' )
	def 	UpperCamelCase_  ( self					:				Optional[int] ):







							'''simple docstring'''
							pass





	@unittest.skip('Esm does not support embedding resizing' )
	def 	UpperCamelCase_  ( self					:				List[Any] ):







							'''simple docstring'''
							pass





	@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.' )
	def 	UpperCamelCase_  ( self					:				Any ):







							'''simple docstring'''
							pass






@require_torch
class 							lowercase__	(							lowercase   ):





	@slow
	def 	UpperCamelCase_  ( self					:				str ):







							'''simple docstring'''
							with torch.no_grad():
													_UpperCamelCase					:		Any 				=   EsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D' )
													model.eval()
													_UpperCamelCase					:		int 				=   torch.tensor([[0, 1, 2, 3, 4, 5]] )
													_UpperCamelCase					:		List[str] 				=   model(lowerCamelCase__ )[0]

													_UpperCamelCase					:		List[Any] 				=   33

													_UpperCamelCase					:		Dict 				=   torch.Size((1, 6, vocab_size) )
													self.assertEqual(output.shape ,lowerCamelCase__ )

													_UpperCamelCase					:		Optional[int] 				=   torch.tensor(
													    [[[8.9_2_1_5, -1_0.5_8_9_8, -6.4_6_7_1], [-6.3_9_6_7, -1_3.9_1_1_4, -1.1_2_1_2], [-7.7_8_1_2, -1_3.9_5_1_6, -3.7_4_0_6]]] )
													self.assertTrue(torch.allclose(output[:, :3, :3] ,lowerCamelCase__ ,atol=1E-4 ) )





	@slow
	def 	UpperCamelCase_  ( self					:				str ):







							'''simple docstring'''
							with torch.no_grad():
													_UpperCamelCase					:		Optional[Any] 				=   EsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D' )
													model.eval()

													_UpperCamelCase					:		str 				=   torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]] )
													_UpperCamelCase					:		List[Any] 				=   model(lowerCamelCase__ )[0]
													# compare the actual values for a slice.
													_UpperCamelCase					:		Dict 				=   torch.tensor(
													    [[[0.1_4_4_4, 0.5_4_1_3, 0.3_2_4_8], [0.3_0_3_4, 0.0_0_5_3, 0.3_1_0_8], [0.3_2_2_8, -0.2_4_9_9, 0.3_4_1_5]]] )
													self.assertTrue(torch.allclose(output[:, :3, :3] ,lowerCamelCase__ ,atol=1E-4 ) )




from typing import List, Optional, Tuple, Union

import torch

from ...models import UNetaDModel
from ...schedulers import ScoreSdeVeScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput







class      A__							(     __magic_name__	):
				lowercase			    =  42
				lowercase			    =  42






				def __init__(						self					:      Any    ,     a					:      UNetaDModel    ,     a					:      ScoreSdeVeScheduler   ):






							'''simple docstring'''







							super().__init__()
							self.register_modules(unet=a    ,     scheduler=a   )






				@torch.no_grad()
				def __call__(						self					:      List[str]    ,     a					:      int = 1    ,     a					:      int = 2_000    ,     a					:      Optional[Union[torch.Generator, List[torch.Generator]]] = None    ,     a					:      Optional[str] = "pil"    ,     a					:      bool = True    ,     **a					:      int    ,     ):






							'''simple docstring'''







							lowerCAmelCase__      :      List[str]         =		self.unet.config.sample_size
							lowerCAmelCase__      :      Union[str, Any]         =		(batch_size, 3, img_size, img_size)

							lowerCAmelCase__      :      Tuple         =		self.unet

							lowerCAmelCase__      :      Optional[Any]         =		randn_tensor(a    ,     generator=a   ) * self.scheduler.init_noise_sigma
							lowerCAmelCase__      :      int         =		sample.to(self.device   )

							self.scheduler.set_timesteps(a   )
							self.scheduler.set_sigmas(a   )

							for i, t in enumerate(self.progress_bar(self.scheduler.timesteps   )   ):
										lowerCAmelCase__      :      Dict         =		self.scheduler.sigmas[i] * torch.ones(shape[0]    ,     device=self.device   )

										# correction step
										for _ in range(self.scheduler.config.correct_steps   ):
													lowerCAmelCase__      :      Optional[Any]         =		self.unet(a    ,     a   ).sample
													lowerCAmelCase__      :      Dict         =		self.scheduler.step_correct(a    ,     a    ,     generator=a   ).prev_sample

										# prediction step
										lowerCAmelCase__      :      Optional[int]         =		model(a    ,     a   ).sample
										lowerCAmelCase__      :      Optional[Any]         =		self.scheduler.step_pred(a    ,     a    ,     a    ,     generator=a   )

										lowerCAmelCase__      ,				lowerCAmelCase__      :      str         =		output.prev_sample, output.prev_sample_mean

							lowerCAmelCase__      :      Any         =		sample_mean.clamp(0    ,     1   )
							lowerCAmelCase__      :      List[str]         =		sample.cpu().permute(0    ,     2    ,     3    ,     1   ).numpy()
							if output_type == "pil":
										lowerCAmelCase__      :      int         =		self.numpy_to_pil(a   )

							if not return_dict:
										return (sample,)

							return ImagePipelineOutput(images=a   )