


"""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 			lowerCamelCase_		:





				"""simple docstring"""







				def __init__(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__=13       ,		UpperCAmelCase__=7       ,		UpperCAmelCase__=False       ,		UpperCAmelCase__=True       ,		UpperCAmelCase__=False       ,		UpperCAmelCase__=True       ,		UpperCAmelCase__=33       ,		UpperCAmelCase__=32       ,		UpperCAmelCase__=5       ,		UpperCAmelCase__=4       ,		UpperCAmelCase__=37       ,		UpperCAmelCase__="gelu"       ,		UpperCAmelCase__=0.1       ,		UpperCAmelCase__=0.1       ,		UpperCAmelCase__=512       ,		UpperCAmelCase__=16       ,		UpperCAmelCase__=2       ,		UpperCAmelCase__=0.02       ,		UpperCAmelCase__=3       ,		UpperCAmelCase__=4       ,		UpperCAmelCase__=None       ,		):
								SCREAMING_SNAKE_CASE__     							=				parent
								SCREAMING_SNAKE_CASE__     							=				batch_size
								SCREAMING_SNAKE_CASE__     							=				seq_length
								SCREAMING_SNAKE_CASE__     							=				is_training
								SCREAMING_SNAKE_CASE__     							=				use_input_mask
								SCREAMING_SNAKE_CASE__     							=				use_token_type_ids
								SCREAMING_SNAKE_CASE__     							=				use_labels
								SCREAMING_SNAKE_CASE__     							=				vocab_size
								SCREAMING_SNAKE_CASE__     							=				hidden_size
								SCREAMING_SNAKE_CASE__     							=				num_hidden_layers
								SCREAMING_SNAKE_CASE__     							=				num_attention_heads
								SCREAMING_SNAKE_CASE__     							=				intermediate_size
								SCREAMING_SNAKE_CASE__     							=				hidden_act
								SCREAMING_SNAKE_CASE__     							=				hidden_dropout_prob
								SCREAMING_SNAKE_CASE__     							=				attention_probs_dropout_prob
								SCREAMING_SNAKE_CASE__     							=				max_position_embeddings
								SCREAMING_SNAKE_CASE__     							=				type_vocab_size
								SCREAMING_SNAKE_CASE__     							=				type_sequence_label_size
								SCREAMING_SNAKE_CASE__     							=				initializer_range
								SCREAMING_SNAKE_CASE__     							=				num_labels
								SCREAMING_SNAKE_CASE__     							=				num_choices
								SCREAMING_SNAKE_CASE__     							=				scope



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				ids_tensor([self.batch_size, self.seq_length]       ,		self.vocab_size     )

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

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

								SCREAMING_SNAKE_CASE__     							=				self.get_config()

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



				def   lowerCAmelCase__		(						self     ):
								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   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):
								SCREAMING_SNAKE_CASE__     							=				EsmModel(config=UpperCAmelCase__     )
								model.to(UpperCAmelCase__     )
								model.eval()
								SCREAMING_SNAKE_CASE__     							=				model(UpperCAmelCase__       ,		attention_mask=UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				model(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				model(UpperCAmelCase__     )

								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   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):
								SCREAMING_SNAKE_CASE__     							=				EsmForMaskedLM(config=UpperCAmelCase__     )
								model.to(UpperCAmelCase__     )
								model.eval()
								SCREAMING_SNAKE_CASE__     							=				model(UpperCAmelCase__       ,		attention_mask=UpperCAmelCase__       ,		labels=UpperCAmelCase__     )
								self.parent.assertEqual(result.logits.shape       ,		(self.batch_size, self.seq_length, self.vocab_size)     )



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):
								SCREAMING_SNAKE_CASE__     							=				self.num_labels
								SCREAMING_SNAKE_CASE__     							=				EsmForTokenClassification(config=UpperCAmelCase__     )
								model.to(UpperCAmelCase__     )
								model.eval()
								SCREAMING_SNAKE_CASE__     							=				model(UpperCAmelCase__       ,		attention_mask=UpperCAmelCase__       ,		labels=UpperCAmelCase__     )
								self.parent.assertEqual(result.logits.shape       ,		(self.batch_size, self.seq_length, self.num_labels)     )






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


@require_torch
class 			lowerCamelCase_		(			lowercase       ,  lowercase       ,  unittest.TestCase   ):





				"""simple docstring"""







				_lowerCAmelCase      :      List[Any]      							=					False

				_lowerCAmelCase      :      List[Any]      							=					(
				    (
				        EsmForMaskedLM,
				        EsmModel,
				        EsmForSequenceClassification,
				        EsmForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				_lowerCAmelCase      :      Optional[int]      							=					()
				_lowerCAmelCase      :      int      							=					(
				    {
				        "feature-extraction": EsmModel,
				        "fill-mask": EsmForMaskedLM,
				        "text-classification": EsmForSequenceClassification,
				        "token-classification": EsmForTokenClassification,
				        "zero-shot": EsmForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				_lowerCAmelCase      :      int      							=					True



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				EsmModelTester(self     )
								SCREAMING_SNAKE_CASE__     							=				ConfigTester(self       ,		config_class=UpperCAmelCase__       ,		hidden_size=37     )



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



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.model_tester.prepare_config_and_inputs()
								for type in ["absolute", "relative_key", "relative_key_query"]:
												SCREAMING_SNAKE_CASE__     							=				type
												self.model_tester.create_and_check_model(*UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*UpperCAmelCase__     )



				@slow
				def   lowerCAmelCase__		(						self     ):
								for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												SCREAMING_SNAKE_CASE__     							=				EsmModel.from_pretrained(UpperCAmelCase__     )
												self.assertIsNotNone(UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.model_tester.prepare_config_and_inputs()[0]
								SCREAMING_SNAKE_CASE__     							=				EsmEmbeddings(config=UpperCAmelCase__     )

								SCREAMING_SNAKE_CASE__     							=				torch.as_tensor([[12, 31, 13, model.padding_idx]]     )
								SCREAMING_SNAKE_CASE__     							=				torch.as_tensor(
								    [
								        [
								            0 + model.padding_idx + 1,
								            1 + model.padding_idx + 1,
								            2 + model.padding_idx + 1,
								            model.padding_idx,
								        ]
								    ]     )
								SCREAMING_SNAKE_CASE__     							=				create_position_ids_from_input_ids(UpperCAmelCase__       ,		model.padding_idx     )
								self.assertEqual(position_ids.shape       ,		expected_positions.shape     )
								self.assertTrue(torch.all(torch.eq(UpperCAmelCase__       ,		UpperCAmelCase__     )     )     )



				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				self.model_tester.prepare_config_and_inputs()[0]
								SCREAMING_SNAKE_CASE__     							=				EsmEmbeddings(config=UpperCAmelCase__     )

								SCREAMING_SNAKE_CASE__     							=				torch.empty(2       ,		4       ,		30     )
								SCREAMING_SNAKE_CASE__     							=				[
								    0 + embeddings.padding_idx + 1,
								    1 + embeddings.padding_idx + 1,
								    2 + embeddings.padding_idx + 1,
								    3 + embeddings.padding_idx + 1,
								]
								SCREAMING_SNAKE_CASE__     							=				torch.as_tensor([expected_single_positions, expected_single_positions]     )
								SCREAMING_SNAKE_CASE__     							=				embeddings.create_position_ids_from_inputs_embeds(UpperCAmelCase__     )
								self.assertEqual(position_ids.shape       ,		expected_positions.shape     )
								self.assertTrue(torch.all(torch.eq(UpperCAmelCase__       ,		UpperCAmelCase__     )     )     )



				@unittest.skip("Esm does not support embedding resizing"     )
				def   lowerCAmelCase__		(						self     ):
								pass



				@unittest.skip("Esm does not support embedding resizing"     )
				def   lowerCAmelCase__		(						self     ):
								pass






				@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."     )
				def   lowerCAmelCase__		(						self     ):
								pass


@require_torch
class 			lowerCamelCase_		(			lowercase   ):





				"""simple docstring"""







				@slow
				def   lowerCAmelCase__		(						self     ):
								with torch.no_grad():
												SCREAMING_SNAKE_CASE__     							=				EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"     )
												model.eval()
												SCREAMING_SNAKE_CASE__     							=				torch.tensor([[0, 1, 2, 3, 4, 5]]     )
												SCREAMING_SNAKE_CASE__     							=				model(UpperCAmelCase__     )[0]

												SCREAMING_SNAKE_CASE__     							=				33

												SCREAMING_SNAKE_CASE__     							=				torch.Size((1, 6, vocab_size)     )
												self.assertEqual(output.shape       ,		UpperCAmelCase__     )

												SCREAMING_SNAKE_CASE__     							=				torch.tensor(
												    [[[8.9_215, -10.5_898, -6.4_671], [-6.3_967, -13.9_114, -1.1_212], [-7.7_812, -13.9_516, -3.7_406]]]     )
												self.assertTrue(torch.allclose(output[:, :3, :3]       ,		UpperCAmelCase__       ,		atol=1e-4     )     )






				@slow
				def   lowerCAmelCase__		(						self     ):
								with torch.no_grad():
												SCREAMING_SNAKE_CASE__     							=				EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"     )
												model.eval()

												SCREAMING_SNAKE_CASE__     							=				torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]     )
												SCREAMING_SNAKE_CASE__     							=				model(UpperCAmelCase__     )[0]
												# compare the actual values for a slice.
												SCREAMING_SNAKE_CASE__     							=				torch.tensor(
												    [[[0.1_444, 0.5_413, 0.3_248], [0.3_034, 0.0_053, 0.3_108], [0.3_228, -0.2_499, 0.3_415]]]     )
												self.assertTrue(torch.allclose(output[:, :3, :3]       ,		UpperCAmelCase__       ,		atol=1e-4     )     )




"""simple docstring"""






# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch

import math
from typing import Union

import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import randn_tensor
from .scheduling_utils import SchedulerMixin


class 			lowerCamelCase_		(			lowercase       ,  lowercase   ):





				"""simple docstring"""







				_lowerCAmelCase      :      List[Any]      							=					1



				@register_to_config
				def __init__(						self       ,		UpperCAmelCase__=2000       ,		UpperCAmelCase__=0.1       ,		UpperCAmelCase__=20       ,		UpperCAmelCase__=1e-3     ):
								SCREAMING_SNAKE_CASE__     							=				None
								SCREAMING_SNAKE_CASE__     							=				None
								SCREAMING_SNAKE_CASE__     							=				None



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__ = None     ):
								SCREAMING_SNAKE_CASE__     							=				torch.linspace(1       ,		self.config.sampling_eps       ,		UpperCAmelCase__       ,		device=UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__=None     ):
								if self.timesteps is None:
												raise ValueError(
												    "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler"     )

								# TODO(Patrick) better comments + non-PyTorch
								# postprocess model score
								SCREAMING_SNAKE_CASE__     							=				(
								    -0.25 * t**2 * (self.config.beta_max - self.config.beta_min) - 0.5 * t * self.config.beta_min
								)
								SCREAMING_SNAKE_CASE__     							=				torch.sqrt(1.0 - torch.exp(2.0 * log_mean_coeff     )     )
								SCREAMING_SNAKE_CASE__     							=				std.flatten()
								while len(std.shape     ) < len(score.shape     ):
												SCREAMING_SNAKE_CASE__     							=				std.unsqueeze(-1     )
								SCREAMING_SNAKE_CASE__     							=				-score / std

								# compute
								SCREAMING_SNAKE_CASE__     							=				-1.0 / len(self.timesteps     )

								SCREAMING_SNAKE_CASE__     							=				self.config.beta_min + t * (self.config.beta_max - self.config.beta_min)
								SCREAMING_SNAKE_CASE__     							=				beta_t.flatten()
								while len(beta_t.shape     ) < len(x.shape     ):
												SCREAMING_SNAKE_CASE__     							=				beta_t.unsqueeze(-1     )
								SCREAMING_SNAKE_CASE__     							=				-0.5 * beta_t * x

								SCREAMING_SNAKE_CASE__     							=				torch.sqrt(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				drift - diffusion**2 * score
								SCREAMING_SNAKE_CASE__     							=				x + drift * dt

								# add noise
								SCREAMING_SNAKE_CASE__     							=				randn_tensor(x.shape       ,		layout=x.layout       ,		generator=UpperCAmelCase__       ,		device=x.device       ,		dtype=x.dtype     )
								SCREAMING_SNAKE_CASE__     							=				x_mean + diffusion * math.sqrt(-dt     ) * noise

								return x, x_mean






				def __len__(						self     ):
								return self.config.num_train_timesteps
