






'''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 UpperCAmelCase_       :




      """simple docstring"""




      def __init__(			self      :		Dict  ,   UpperCAmelCase      :		Any  ,   )	->	Dict:


       '''simple docstring'''

       lowercase       : Optional[Any]						=parent
       lowercase       : Union[str, Any]						=13
       lowercase       : Dict						=7
       lowercase       : Optional[Any]						=True
       lowercase       : Dict						=True
       lowercase       : Optional[int]						=True
       lowercase       : int						=99
       lowercase       : int						=32
       lowercase       : int						=2
       lowercase       : Union[str, Any]						=4
       lowercase       : Dict						=37
       lowercase       : Optional[int]						='''gelu'''
       lowercase       : int						=0.1
       lowercase       : Optional[int]						=0.1
       lowercase       : int						=512
       lowercase       : List[Any]						=16
       lowercase       : Any						=2
       lowercase       : int						=0.0_2
       lowercase       : Any						=3
       lowercase       : Dict						=4
       lowercase       : Union[str, Any]						=None






      def 					A__   (			self      :		Dict		)	->	Any:


       '''simple docstring'''

       lowercase       : int						=ids_tensor([self.batch_size, self.seq_length]  ,   self.vocab_size		)

       lowercase       : Dict						=None
       if self.use_input_mask:
        lowercase       : List[str]						=random_attention_mask([self.batch_size, self.seq_length]		)

       lowercase       : Optional[Any]						=None
       lowercase       : int						=None
       lowercase       : Optional[Any]						=None
       if self.use_labels:
        lowercase       : str						=ids_tensor([self.batch_size]  ,   self.type_sequence_label_size		)
        lowercase       : Tuple						=ids_tensor([self.batch_size, self.seq_length]  ,   self.num_labels		)
        lowercase       : Tuple						=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 					A__   (			self      :		int		)	->	List[str]:


       '''simple docstring'''

       (
           (
           lowercase       
       )	,		(
           lowercase       
       )	,		(
           lowercase       
       )	,		(
           lowercase       
       )	,		(
           lowercase       
       )	,		(
           lowercase       
       )	,		       
       )       : int						=self.prepare_config_and_inputs()

       lowercase       : Union[str, Any]						=True
       lowercase       : Optional[int]						=floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
       lowercase       : Tuple						=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 					A__   (			self      :		List[Any]  ,   UpperCAmelCase      :		List[Any]  ,   UpperCAmelCase      :		Union[str, Any]  ,   UpperCAmelCase      :		int  ,   UpperCAmelCase      :		str  ,   UpperCAmelCase      :		Tuple  ,   UpperCAmelCase      :		Optional[Any]		)	->	List[str]:


       '''simple docstring'''

       lowercase       : List[str]						=TFEsmModel(config=_UpperCAmelCase		)
       lowercase       : str						={'''input_ids''': input_ids, '''attention_mask''': input_mask}
       lowercase       : Tuple						=model(_UpperCAmelCase		)

       lowercase       : List[Any]						=[input_ids, input_mask]
       lowercase       : Dict						=model(_UpperCAmelCase		)

       lowercase       : Optional[int]						=model(_UpperCAmelCase		)

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






      def 					A__   (			self      :		List[Any]  ,   UpperCAmelCase      :		Dict  ,   UpperCAmelCase      :		str  ,   UpperCAmelCase      :		Union[str, Any]  ,   UpperCAmelCase      :		Optional[int]  ,   UpperCAmelCase      :		Union[str, Any]  ,   UpperCAmelCase      :		Optional[int]  ,   UpperCAmelCase      :		List[Any]  ,   UpperCAmelCase      :		Dict  ,   )	->	Any:


       '''simple docstring'''

       lowercase       : int						=True

       lowercase       : Any						=TFEsmModel(config=_UpperCAmelCase		)
       lowercase       : int						={
           '''input_ids''': input_ids,
           '''attention_mask''': input_mask,
           '''encoder_hidden_states''': encoder_hidden_states,
           '''encoder_attention_mask''': encoder_attention_mask,
       }
       lowercase       : str						=model(_UpperCAmelCase		)

       lowercase       : str						=[input_ids, input_mask]
       lowercase       : Optional[Any]						=model(_UpperCAmelCase  ,   encoder_hidden_states=_UpperCAmelCase		)

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

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






      def 					A__   (			self      :		Optional[int]  ,   UpperCAmelCase      :		Tuple  ,   UpperCAmelCase      :		Optional[int]  ,   UpperCAmelCase      :		Any  ,   UpperCAmelCase      :		Dict  ,   UpperCAmelCase      :		Dict  ,   UpperCAmelCase      :		Optional[int]		)	->	Optional[int]:


       '''simple docstring'''

       lowercase       : List[str]						=TFEsmForMaskedLM(config=_UpperCAmelCase		)
       lowercase       : Any						=model([input_ids, input_mask]		)
       self.parent.assertEqual(result.logits.shape  ,   (self.batch_size, self.seq_length, self.vocab_size)		)






      def 					A__   (			self      :		Optional[int]  ,   UpperCAmelCase      :		List[Any]  ,   UpperCAmelCase      :		List[Any]  ,   UpperCAmelCase      :		List[Any]  ,   UpperCAmelCase      :		Tuple  ,   UpperCAmelCase      :		List[str]  ,   UpperCAmelCase      :		Tuple		)	->	str:


       '''simple docstring'''

       lowercase       : Dict						=self.num_labels
       lowercase       : Union[str, Any]						=TFEsmForTokenClassification(config=_UpperCAmelCase		)
       lowercase       : int						={'''input_ids''': input_ids, '''attention_mask''': input_mask}
       lowercase       : List[str]						=model(_UpperCAmelCase		)
       self.parent.assertEqual(result.logits.shape  ,   (self.batch_size, self.seq_length, self.num_labels)		)






      def 					A__   (			self      :		Optional[int]		)	->	Dict:


       '''simple docstring'''

       lowercase       : int						=self.prepare_config_and_inputs()
       (
           (
           lowercase       
       )	,		(
           lowercase       
       )	,		(
           lowercase       
       )	,		(
           lowercase       
       )	,		(
           lowercase       
       )	,		(
           lowercase       
       )	,		       
       )       : Tuple						=config_and_inputs
       lowercase       : Tuple						={'''input_ids''': input_ids, '''attention_mask''': input_mask}
       return config, inputs_dict





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




      """simple docstring"""




      UpperCamelCase_             =      (
          (
              TFEsmModel,
              TFEsmForMaskedLM,
              TFEsmForSequenceClassification,
              TFEsmForTokenClassification,
          )
          if is_tf_available()
          else ()
      )
      UpperCamelCase_             =      (
          {
              '''feature-extraction''': TFEsmModel,
              '''fill-mask''': TFEsmForMaskedLM,
              '''text-classification''': TFEsmForSequenceClassification,
              '''token-classification''': TFEsmForTokenClassification,
              '''zero-shot''': TFEsmForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      UpperCamelCase_             =      False
      UpperCamelCase_             =      False






      def 					A__   (			self      :		Union[str, Any]		)	->	Optional[int]:


       '''simple docstring'''

       lowercase       : str						=TFEsmModelTester(self		)
       lowercase       : Optional[Any]						=ConfigTester(self  ,   config_class=_UpperCAmelCase  ,   hidden_size=37		)






      def 					A__   (			self      :		Union[str, Any]		)	->	List[str]:


       '''simple docstring'''

       self.config_tester.run_common_tests()






      def 					A__   (			self      :		int		)	->	List[Any]:


       '''simple docstring'''

       lowercase       : Tuple						=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_UpperCAmelCase		)






      def 					A__   (			self      :		List[Any]		)	->	Optional[int]:


       '''simple docstring'''

       lowercase       : Union[str, Any]						=self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_model_as_decoder(*_UpperCAmelCase		)






      def 					A__   (			self      :		List[Any]		)	->	int:


       '''simple docstring'''

       lowercase       : int						=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*_UpperCAmelCase		)






      def 					A__   (			self      :		Dict		)	->	Tuple:


       '''simple docstring'''

       lowercase       : Union[str, Any]						=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*_UpperCAmelCase		)






      @slow
      def 					A__   (			self      :		List[str]		)	->	Optional[Any]:


       '''simple docstring'''

       for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        lowercase       : Optional[int]						=TFEsmModel.from_pretrained(_UpperCAmelCase		)
        self.assertIsNotNone(_UpperCAmelCase		)






      @unittest.skip('''Protein models do not support embedding resizing.'''		)
      def 					A__   (			self      :		str		)	->	Dict:


       '''simple docstring'''

       pass






      @unittest.skip('''Protein models do not support embedding resizing.'''		)
      def 					A__   (			self      :		int		)	->	List[Any]:


       '''simple docstring'''

       pass






      def 					A__   (			self      :		Dict		)	->	int:


       '''simple docstring'''

       lowercase	,		lowercase       : List[str]						=self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        lowercase       : Any						=model_class(_UpperCAmelCase		)
        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       : Optional[int]						=model.get_bias()
         assert isinstance(_UpperCAmelCase  ,   _UpperCAmelCase		)
         for k, v in name.items():
          assert isinstance(_UpperCAmelCase  ,   tf.Variable		)
        else:
         lowercase       : Union[str, Any]						=model.get_output_embeddings()
         assert x is None
         lowercase       : Dict						=model.get_bias()
         assert name is None





@require_tf
class UpperCAmelCase_       (     unittest.TestCase					):




      """simple docstring"""




      @slow
      def 					A__   (			self      :		List[str]		)	->	Tuple:


       '''simple docstring'''

       lowercase       : Tuple						=TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''		)

       lowercase       : int						=tf.constant([[0, 1, 2, 3, 4, 5]]		)
       lowercase       : Union[str, Any]						=model(_UpperCAmelCase		)[0]
       lowercase       : str						=[1, 6, 33]
       self.assertEqual(list(output.numpy().shape		)  ,   _UpperCAmelCase		)
       # compare the actual values for a slice.
       lowercase       : Dict						=tf.constant(
           [
               [
                   [8.9_2_1_5_1_8, -1_0.5_8_9_8_1_4, -6.4_6_7_1_3_0_7],
                   [-6.3_9_6_7_1_5_6, -1_3.9_1_1_3_7_7, -1.1_2_1_1_9_1_5],
                   [-7.7_8_1_2_4_7, -1_3.9_5_1_5_5_7, -3.7_4_0_5_9_2],
               ]
           ]		)
       self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()  ,   expected_slice.numpy()  ,   atol=1e-2		)		)






      @slow
      def 					A__   (			self      :		int		)	->	List[Any]:


       '''simple docstring'''

       lowercase       : Optional[Any]						=TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''		)

       lowercase       : int						=tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]		)
       lowercase       : List[Any]						=model(_UpperCAmelCase		)[0]
       # compare the actual values for a slice.
       lowercase       : List[Any]						=tf.constant(
           [
               [
                   [0.1_4_4_4_3_0_9_2, 0.5_4_1_2_5_3_2_7, 0.3_2_4_7_7_3_9],
                   [0.3_0_3_4_0_4_8_4, 0.0_0_5_2_6_6_7_6, 0.3_1_0_7_7_7_2_2],
                   [0.3_2_2_7_8_0_4_3, -0.2_4_9_8_7_0_9_6, 0.3_4_1_4_6_2_8],
               ]
           ]		)
       self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()  ,   expected_slice.numpy()  ,   atol=1e-4		)		)









"""simple docstring"""

import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest





class 							lowercase_			(	__lowerCAmelCase				):



 '''simple docstring'''
 UpperCAmelCase       :	Optional[int]     =      (KDPMaDiscreteScheduler,)
 UpperCAmelCase       :	Any     =      10

 def      lowerCAmelCase_				(					self			:			Dict       ,				**_UpperCAmelCase			:			Optional[Any]						):
       _A     =  {
           'num_train_timesteps': 1_100,
           'beta_start': 0.0001,
           'beta_end': 0.02,
           'beta_schedule': 'linear',
       }

       config.update(**_UpperCAmelCase						)
       return config

 def      lowerCAmelCase_				(					self			:			Any						):
       for timesteps in [10, 50, 100, 1_000]:
             self.check_over_configs(num_train_timesteps=_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			Dict						):
       for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001]       ,				[0.0002, 0.002, 0.02]						):
             self.check_over_configs(beta_start=_UpperCAmelCase       ,				beta_end=_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			Tuple						):
       for schedule in ["linear", "scaled_linear"]:
             self.check_over_configs(beta_schedule=_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			Optional[int]						):
       for prediction_type in ["epsilon", "v_prediction"]:
             self.check_over_configs(prediction_type=_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			Optional[int]						):
       _A     =  self.scheduler_classes[0]
       _A     =  self.get_scheduler_config(prediction_type='v_prediction'						)
       _A     =  scheduler_class(**_UpperCAmelCase						)

       scheduler.set_timesteps(self.num_inference_steps						)

       _A     =  self.dummy_model()
       _A     =  self.dummy_sample_deter * scheduler.init_noise_sigma
       _A     =  sample.to(_UpperCAmelCase						)

       for i, t in enumerate(scheduler.timesteps						):
             _A     =  scheduler.scale_model_input(_UpperCAmelCase       ,				_UpperCAmelCase						)

             _A     =  model(_UpperCAmelCase       ,				_UpperCAmelCase						)

             _A     =  scheduler.step(_UpperCAmelCase       ,				_UpperCAmelCase       ,				_UpperCAmelCase						)
             _A     =  output.prev_sample

       _A     =  torch.sum(torch.abs(_UpperCAmelCase						)						)
       _A     =  torch.mean(torch.abs(_UpperCAmelCase						)						)

       if torch_device in ["cpu", "mps"]:
             assert abs(result_sum.item() - 4.6_9_3_4E-0_7						) < 1E-2
             assert abs(result_mean.item() - 6.1_1_1_2E-1_0						) < 1E-3
       else:
             # CUDA
             assert abs(result_sum.item() - 4.6_9_3_4_2_8_6_5_0_1_7_0_9_7_2E-0_7						) < 1E-2
             assert abs(result_mean.item() - 0.0002						) < 1E-3

 def      lowerCAmelCase_				(					self			:			Optional[Any]						):
       if torch_device == "mps":
             return
       _A     =  self.scheduler_classes[0]
       _A     =  self.get_scheduler_config()
       _A     =  scheduler_class(**_UpperCAmelCase						)

       scheduler.set_timesteps(self.num_inference_steps						)

       _A     =  self.dummy_model()
       _A     =  self.dummy_sample_deter * scheduler.init_noise_sigma
       _A     =  sample.to(_UpperCAmelCase						)

       for i, t in enumerate(scheduler.timesteps						):
             _A     =  scheduler.scale_model_input(_UpperCAmelCase       ,				_UpperCAmelCase						)

             _A     =  model(_UpperCAmelCase       ,				_UpperCAmelCase						)

             _A     =  scheduler.step(_UpperCAmelCase       ,				_UpperCAmelCase       ,				_UpperCAmelCase						)
             _A     =  output.prev_sample

       _A     =  torch.sum(torch.abs(_UpperCAmelCase						)						)
       _A     =  torch.mean(torch.abs(_UpperCAmelCase						)						)

       if torch_device in ["cpu", "mps"]:
             assert abs(result_sum.item() - 20.4125						) < 1E-2
             assert abs(result_mean.item() - 0.0266						) < 1E-3
       else:
             # CUDA
             assert abs(result_sum.item() - 20.4125						) < 1E-2
             assert abs(result_mean.item() - 0.0266						) < 1E-3



 def      lowerCAmelCase_				(					self			:			Any						):
       if torch_device == "mps":
             return
       _A     =  self.scheduler_classes[0]
       _A     =  self.get_scheduler_config()
       _A     =  scheduler_class(**_UpperCAmelCase						)

       scheduler.set_timesteps(self.num_inference_steps       ,				device=_UpperCAmelCase						)

       _A     =  self.dummy_model()
       _A     =  self.dummy_sample_deter.to(_UpperCAmelCase						) * scheduler.init_noise_sigma

       for t in scheduler.timesteps:
             _A     =  scheduler.scale_model_input(_UpperCAmelCase       ,				_UpperCAmelCase						)

             _A     =  model(_UpperCAmelCase       ,				_UpperCAmelCase						)

             _A     =  scheduler.step(_UpperCAmelCase       ,				_UpperCAmelCase       ,				_UpperCAmelCase						)
             _A     =  output.prev_sample

       _A     =  torch.sum(torch.abs(_UpperCAmelCase						)						)
       _A     =  torch.mean(torch.abs(_UpperCAmelCase						)						)

       if str(_UpperCAmelCase						).startswith('cpu'						):
             # The following sum varies between 148 and 156 on mps. Why?
             assert abs(result_sum.item() - 20.4125						) < 1E-2
             assert abs(result_mean.item() - 0.0266						) < 1E-3
       else:
             # CUDA
             assert abs(result_sum.item() - 20.4125						) < 1E-2
             assert abs(result_mean.item() - 0.0266						) < 1E-3

