






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


							def __init__(			self      :						Any			,    snake_case      :						str			,    ):
													'''simple docstring'''





													A__						:  Optional[int]		  =  parent
													A__						:  str		  =  13
													A__						:  Optional[int]		  =  7
													A__						:  Optional[Any]		  =  True
													A__						:  int		  =  True
													A__						:  List[str]		  =  True
													A__						:  Tuple		  =  99
													A__						:  Union[str, Any]		  =  32
													A__						:  str		  =  2
													A__						:  Tuple		  =  4
													A__						:  Any		  =  37
													A__						:  int		  =  """gelu"""
													A__						:  Any		  =  0.1
													A__						:  Tuple		  =  0.1
													A__						:  Optional[Any]		  =  512
													A__						:  List[str]		  =  16
													A__						:  List[str]		  =  2
													A__						:  Dict		  =  0.02
													A__						:  Tuple		  =  3
													A__						:  List[Any]		  =  4
													A__						:  Dict		  =  None


							def 						_UpperCamelCase       (			self      :						List[Any]							):
													'''simple docstring'''





													A__						:  Tuple		  =  ids_tensor([self.batch_size, self.seq_length]			,    self.vocab_size							)

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

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

													A__						:  Optional[Any]		  =  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 						_UpperCamelCase       (			self      :						Tuple							):
													'''simple docstring'''





													(
													    (
													    A__
													)  ,  (
													    A__
													)  ,  (
													    A__
													)  ,  (
													    A__
													)  ,  (
													    A__
													)  ,  (
													    A__
													)  ,  
													)						:  Tuple		  =  self.prepare_config_and_inputs()

													A__						:  Dict		  =  True
													A__						:  Tuple		  =  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
													A__						:  List[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 						_UpperCamelCase       (			self      :						Optional[int]			,    snake_case      :						Dict			,    snake_case      :						List[str]			,    snake_case      :						Optional[int]			,    snake_case      :						Any			,    snake_case      :						Any			,    snake_case      :						Optional[int]							):
													'''simple docstring'''





													A__						:  Optional[Any]		  =  TFEsmModel(config=snake_case							)
													A__						:  Optional[int]		  =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
													A__						:  Any		  =  model(snake_case							)

													A__						:  List[Any]		  =  [input_ids, input_mask]
													A__						:  int		  =  model(snake_case							)

													A__						:  Union[str, Any]		  =  model(snake_case							)

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


							def 						_UpperCamelCase       (			self      :						Optional[int]			,    snake_case      :						List[str]			,    snake_case      :						List[Any]			,    snake_case      :						List[Any]			,    snake_case      :						List[str]			,    snake_case      :						List[str]			,    snake_case      :						List[Any]			,    snake_case      :						int			,    snake_case      :						Optional[Any]			,    ):
													'''simple docstring'''





													A__						:  Tuple		  =  True

													A__						:  List[str]		  =  TFEsmModel(config=snake_case							)
													A__						:  Union[str, Any]		  =  {
													    """input_ids""": input_ids,
													    """attention_mask""": input_mask,
													    """encoder_hidden_states""": encoder_hidden_states,
													    """encoder_attention_mask""": encoder_attention_mask,
													}
													A__						:  List[str]		  =  model(snake_case							)

													A__						:  Union[str, Any]		  =  [input_ids, input_mask]
													A__						:  Any		  =  model(snake_case			,    encoder_hidden_states=snake_case							)

													# Also check the case where encoder outputs are not passed
													A__						:  str		  =  model(snake_case			,    attention_mask=snake_case							)

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


							def 						_UpperCamelCase       (			self      :						int			,    snake_case      :						Optional[int]			,    snake_case      :						Dict			,    snake_case      :						Optional[int]			,    snake_case      :						str			,    snake_case      :						Tuple			,    snake_case      :						str							):
													'''simple docstring'''





													A__						:  Optional[int]		  =  TFEsmForMaskedLM(config=snake_case							)
													A__						:  List[Any]		  =  model([input_ids, input_mask]							)
													self.parent.assertEqual(result.logits.shape			,    (self.batch_size, self.seq_length, self.vocab_size)							)


							def 						_UpperCamelCase       (			self      :						int			,    snake_case      :						Any			,    snake_case      :						Optional[int]			,    snake_case      :						Dict			,    snake_case      :						Tuple			,    snake_case      :						Any			,    snake_case      :						Tuple							):
													'''simple docstring'''





													A__						:  Dict		  =  self.num_labels
													A__						:  List[Any]		  =  TFEsmForTokenClassification(config=snake_case							)
													A__						:  List[Any]		  =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
													A__						:  Optional[Any]		  =  model(snake_case							)
													self.parent.assertEqual(result.logits.shape			,    (self.batch_size, self.seq_length, self.num_labels)							)


							def 						_UpperCamelCase       (			self      :						Optional[int]							):
													'''simple docstring'''





													A__						:  Tuple		  =  self.prepare_config_and_inputs()
													(
													    (
													    A__
													)  ,  (
													    A__
													)  ,  (
													    A__
													)  ,  (
													    A__
													)  ,  (
													    A__
													)  ,  (
													    A__
													)  ,  
													)						:  int		  =  config_and_inputs
													A__						:  str		  =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
													return config, inputs_dict






@require_tf
class 							__SCREAMING_SNAKE_CASE						(				UpperCamelCase       ,		UpperCamelCase       ,		unittest.TestCase		):
							snake_case_     =							(
							    (
							        TFEsmModel,
							        TFEsmForMaskedLM,
							        TFEsmForSequenceClassification,
							        TFEsmForTokenClassification,
							    )
							    if is_tf_available()
							    else ()
							)
							snake_case_     =							(
							    {
							        'feature-extraction': TFEsmModel,
							        'fill-mask': TFEsmForMaskedLM,
							        'text-classification': TFEsmForSequenceClassification,
							        'token-classification': TFEsmForTokenClassification,
							        'zero-shot': TFEsmForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							snake_case_     =							False
							snake_case_     =							False


							def 						_UpperCamelCase       (			self      :						Optional[Any]							):
													'''simple docstring'''





													A__						:  Dict		  =  TFEsmModelTester(self							)
													A__						:  int		  =  ConfigTester(self			,    config_class=snake_case			,    hidden_size=37							)


							def 						_UpperCamelCase       (			self      :						Tuple							):
													'''simple docstring'''





													self.config_tester.run_common_tests()


							def 						_UpperCamelCase       (			self      :						int							):
													'''simple docstring'''





													A__						:  List[Any]		  =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*snake_case							)


							def 						_UpperCamelCase       (			self      :						Union[str, Any]							):
													'''simple docstring'''





													A__						:  Optional[int]		  =  self.model_tester.prepare_config_and_inputs_for_decoder()
													self.model_tester.create_and_check_model_as_decoder(*snake_case							)


							def 						_UpperCamelCase       (			self      :						int							):
													'''simple docstring'''





													A__						:  Tuple		  =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_masked_lm(*snake_case							)


							def 						_UpperCamelCase       (			self      :						Tuple							):
													'''simple docstring'''





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


							@slow
							def 						_UpperCamelCase       (			self      :						Optional[int]							):
													'''simple docstring'''





													for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			A__						:  Any		  =  TFEsmModel.from_pretrained(snake_case							)
																			self.assertIsNotNone(snake_case							)


							@unittest.skip("""Protein models do not support embedding resizing."""							)
							def 						_UpperCamelCase       (			self      :						Optional[int]							):
													'''simple docstring'''





													pass


							@unittest.skip("""Protein models do not support embedding resizing."""							)
							def 						_UpperCamelCase       (			self      :						Union[str, Any]							):
													'''simple docstring'''





													pass


							def 						_UpperCamelCase       (			self      :						Optional[int]							):
													'''simple docstring'''





													A__  ,  A__						:  Optional[Any]		  =  self.model_tester.prepare_config_and_inputs_for_common()

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






@require_tf
class 							__SCREAMING_SNAKE_CASE						(				unittest.TestCase		):


							@slow
							def 						_UpperCamelCase       (			self      :						Any							):
													'''simple docstring'''





													A__						:  Any		  =  TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""							)

													A__						:  List[str]		  =  tf.constant([[0, 1, 2, 3, 4, 5]]							)
													A__						:  Any		  =  model(snake_case							)[0]
													A__						:  List[str]		  =  [1, 6, 33]
													self.assertEqual(list(output.numpy().shape							)			,    snake_case							)
													# compare the actual values for a slice.
													A__						:  Union[str, Any]		  =  tf.constant(
													    [
													        [
													            [8.921518, -10.589814, -6.4671307],
													            [-6.3967156, -13.911377, -1.1211915],
													            [-7.781247, -13.951557, -3.740592],
													        ]
													    ]							)
													self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()			,    expected_slice.numpy()			,    atol=1e-2							)							)


							@slow
							def 						_UpperCamelCase       (			self      :						List[Any]							):
													'''simple docstring'''





													A__						:  Optional[Any]		  =  TFEsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""							)

													A__						:  Tuple		  =  tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]							)
													A__						:  Optional[Any]		  =  model(snake_case							)[0]
													# compare the actual values for a slice.
													A__						:  Any		  =  tf.constant(
													    [
													        [
													            [0.14443092, 0.54125327, 0.3247739],
													            [0.30340484, 0.00526676, 0.31077722],
													            [0.32278043, -0.24987096, 0.3414628],
													        ]
													    ]							)
													self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()			,    expected_slice.numpy()			,    atol=1e-4							)							)














"""simple docstring"""



import unittest

import numpy as np
import torch

from diffusers import ScoreSdeVePipeline, ScoreSdeVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()






class 							__SCREAMING_SNAKE_CASE						(				unittest.TestCase		):


							@property
							def 						_UpperCamelCase       (			self      :						Dict							):
													'''simple docstring'''





													torch.manual_seed(0							)
													A__						:  Optional[Any]		  =  UNetaDModel(
													    block_out_channels=(32, 64)			,    layers_per_block=2			,    sample_size=32			,    in_channels=3			,    out_channels=3			,    down_block_types=("""DownBlock2D""", """AttnDownBlock2D""")			,    up_block_types=("""AttnUpBlock2D""", """UpBlock2D""")			,    )
													return model


							def 						_UpperCamelCase       (			self      :						Optional[Any]							):
													'''simple docstring'''





													A__						:  List[Any]		  =  self.dummy_uncond_unet
													A__						:  Dict		  =  ScoreSdeVeScheduler()

													A__						:  str		  =  ScoreSdeVePipeline(unet=snake_case			,    scheduler=snake_case							)
													sde_ve.to(snake_case							)
													sde_ve.set_progress_bar_config(disable=snake_case							)

													A__						:  Optional[int]		  =  torch.manual_seed(0							)
													A__						:  Any		  =  sde_ve(num_inference_steps=2			,    output_type="""numpy"""			,    generator=snake_case							).images

													A__						:  List[str]		  =  torch.manual_seed(0							)
													A__						:  int		  =  sde_ve(num_inference_steps=2			,    output_type="""numpy"""			,    generator=snake_case			,    return_dict=snake_case							)[
													    0
													]

													A__						:  int		  =  image[0, -3:, -3:, -1]
													A__						:  Optional[Any]		  =  image_from_tuple[0, -3:, -3:, -1]

													assert image.shape == (1, 32, 32, 3)
													A__						:  Optional[Any]		  =  np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]							)

													assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2
													assert np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < 1e-2






@slow
@require_torch
class 							__SCREAMING_SNAKE_CASE						(				unittest.TestCase		):


							def 						_UpperCamelCase       (			self      :						Tuple							):
													'''simple docstring'''





													A__						:  Union[str, Any]		  =  """google/ncsnpp-church-256"""
													A__						:  Tuple		  =  UNetaDModel.from_pretrained(snake_case							)

													A__						:  int		  =  ScoreSdeVeScheduler.from_pretrained(snake_case							)

													A__						:  List[Any]		  =  ScoreSdeVePipeline(unet=snake_case			,    scheduler=snake_case							)
													sde_ve.to(snake_case							)
													sde_ve.set_progress_bar_config(disable=snake_case							)

													A__						:  Dict		  =  torch.manual_seed(0							)
													A__						:  Optional[int]		  =  sde_ve(num_inference_steps=10			,    output_type="""numpy"""			,    generator=snake_case							).images

													A__						:  Optional[int]		  =  image[0, -3:, -3:, -1]

													assert image.shape == (1, 256, 256, 3)

													A__						:  Dict		  =  np.array([0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0]							)

													assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2






