



import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class 	__snake_case   ( SCREAMING_SNAKE_CASE							, unittest.TestCase ):
		SCREAMING_SNAKE_CASE__								=	ConsistencyModelPipeline
		SCREAMING_SNAKE_CASE__								=	UNCONDITIONAL_IMAGE_GENERATION_PARAMS
		SCREAMING_SNAKE_CASE__								=	UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

		# Override required_optional_params to remove num_images_per_prompt
		SCREAMING_SNAKE_CASE__								=	frozenset(
		    [
		        'num_inference_steps',
		        'generator',
		        'latents',
		        'output_type',
		        'return_dict',
		        'callback',
		        'callback_steps',
		    ] )
		@property
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  UNetaDModel.from_pretrained(
						    'diffusers/consistency-models-test'    ,subfolder='test_unet'    ,)
						return unet
		@property
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  UNetaDModel.from_pretrained(
						    'diffusers/consistency-models-test'    ,subfolder='test_unet_class_cond'    ,)
						return unet
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_=False						):







						"""simple docstring"""
						if class_cond:
										lowerCAmelCase__		  =  self.dummy_cond_unet
						else:
										lowerCAmelCase__		  =  self.dummy_uncond_unet

						# Default to CM multistep sampler
						lowerCAmelCase__		  =  CMStochasticIterativeScheduler(
						    num_train_timesteps=40    ,sigma_min=0.002    ,sigma_max=80.0    ,)

						lowerCAmelCase__		  =  {
						    'unet': unet,
						    'scheduler': scheduler,
						}

						return components
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_=0						):







						"""simple docstring"""
						if str(a_						).startswith('mps'						):
										lowerCAmelCase__		  =  torch.manual_seed(a_						)
						else:
										lowerCAmelCase__		  =  torch.Generator(device=a_						).manual_seed(a_						)

						lowerCAmelCase__		  =  {
						    'batch_size': 1,
						    'num_inference_steps': None,
						    'timesteps': [22, 0],
						    'generator': generator,
						    'output_type': 'np',
						}

						return inputs
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  'cpu'  # ensure determinism for the device-dependent torch.Generator
						lowerCAmelCase__		  =  self.get_dummy_components()
						lowerCAmelCase__		  =  ConsistencyModelPipeline(**a_						)
						lowerCAmelCase__		  =  pipe.to(a_						)
						pipe.set_progress_bar_config(disable=a_						)

						lowerCAmelCase__		  =  self.get_dummy_inputs(a_						)
						lowerCAmelCase__		  =  pipe(**a_						).images
						assert image.shape == (1, 32, 32, 3)

						lowerCAmelCase__		  =  image[0, -3:, -3:, -1]
						lowerCAmelCase__		  =  np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]						)

						assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-3
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  'cpu'  # ensure determinism for the device-dependent torch.Generator
						lowerCAmelCase__		  =  self.get_dummy_components(class_cond=a_						)
						lowerCAmelCase__		  =  ConsistencyModelPipeline(**a_						)
						lowerCAmelCase__		  =  pipe.to(a_						)
						pipe.set_progress_bar_config(disable=a_						)

						lowerCAmelCase__		  =  self.get_dummy_inputs(a_						)
						lowerCAmelCase__		  =  0
						lowerCAmelCase__		  =  pipe(**a_						).images
						assert image.shape == (1, 32, 32, 3)

						lowerCAmelCase__		  =  image[0, -3:, -3:, -1]
						lowerCAmelCase__		  =  np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]						)

						assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-3
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  'cpu'  # ensure determinism for the device-dependent torch.Generator
						lowerCAmelCase__		  =  self.get_dummy_components()
						lowerCAmelCase__		  =  ConsistencyModelPipeline(**a_						)
						lowerCAmelCase__		  =  pipe.to(a_						)
						pipe.set_progress_bar_config(disable=a_						)

						lowerCAmelCase__		  =  self.get_dummy_inputs(a_						)
						lowerCAmelCase__		  =  1
						lowerCAmelCase__		  =  None
						lowerCAmelCase__		  =  pipe(**a_						).images
						assert image.shape == (1, 32, 32, 3)

						lowerCAmelCase__		  =  image[0, -3:, -3:, -1]
						lowerCAmelCase__		  =  np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]						)

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



		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  'cpu'  # ensure determinism for the device-dependent torch.Generator
						lowerCAmelCase__		  =  self.get_dummy_components(class_cond=a_						)
						lowerCAmelCase__		  =  ConsistencyModelPipeline(**a_						)
						lowerCAmelCase__		  =  pipe.to(a_						)
						pipe.set_progress_bar_config(disable=a_						)

						lowerCAmelCase__		  =  self.get_dummy_inputs(a_						)
						lowerCAmelCase__		  =  1
						lowerCAmelCase__		  =  None
						lowerCAmelCase__		  =  0
						lowerCAmelCase__		  =  pipe(**a_						).images
						assert image.shape == (1, 32, 32, 3)

						lowerCAmelCase__		  =  image[0, -3:, -3:, -1]
						lowerCAmelCase__		  =  np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]						)

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



@slow
@require_torch_gpu
class 	__snake_case   ( unittest.TestCase ):
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_=0    ,a_=False    ,a_="cpu"    ,a_=torch.floataa    ,a_=(1, 3, 64, 64)						):







						"""simple docstring"""
						lowerCAmelCase__		  =  torch.manual_seed(a_						)

						lowerCAmelCase__		  =  {
						    'num_inference_steps': None,
						    'timesteps': [22, 0],
						    'class_labels': 0,
						    'generator': generator,
						    'output_type': 'np',
						}

						if get_fixed_latents:
										lowerCAmelCase__		  =  self.get_fixed_latents(seed=a_    ,device=a_    ,dtype=a_    ,shape=a_						)
										lowerCAmelCase__		  =  latents

						return inputs
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_=0    ,a_="cpu"    ,a_=torch.floataa    ,a_=(1, 3, 64, 64)						):







						"""simple docstring"""
						if type(a_						) == str:
										lowerCAmelCase__		  =  torch.device(a_						)
						lowerCAmelCase__		  =  torch.Generator(device=a_						).manual_seed(a_						)
						lowerCAmelCase__		  =  randn_tensor(a_    ,generator=a_    ,device=a_    ,dtype=a_						)
						return latents
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  UNetaDModel.from_pretrained('diffusers/consistency_models'    ,subfolder='diffusers_cd_imagenet64_l2'						)
						lowerCAmelCase__		  =  CMStochasticIterativeScheduler(
						    num_train_timesteps=40    ,sigma_min=0.002    ,sigma_max=80.0    ,)
						lowerCAmelCase__		  =  ConsistencyModelPipeline(unet=a_    ,scheduler=a_						)
						pipe.to(torch_device=a_						)
						pipe.set_progress_bar_config(disable=a_						)

						lowerCAmelCase__		  =  self.get_inputs()
						lowerCAmelCase__		  =  pipe(**a_						).images
						assert image.shape == (1, 64, 64, 3)

						lowerCAmelCase__		  =  image[0, -3:, -3:, -1]

						lowerCAmelCase__		  =  np.array([0.0888, 0.0881, 0.0666, 0.0479, 0.0292, 0.0195, 0.0201, 0.0163, 0.0254]						)

						assert np.abs(image_slice.flatten() - expected_slice						).max() < 2e-2
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  UNetaDModel.from_pretrained('diffusers/consistency_models'    ,subfolder='diffusers_cd_imagenet64_l2'						)
						lowerCAmelCase__		  =  CMStochasticIterativeScheduler(
						    num_train_timesteps=40    ,sigma_min=0.002    ,sigma_max=80.0    ,)
						lowerCAmelCase__		  =  ConsistencyModelPipeline(unet=a_    ,scheduler=a_						)
						pipe.to(torch_device=a_						)
						pipe.set_progress_bar_config(disable=a_						)

						lowerCAmelCase__		  =  self.get_inputs()
						lowerCAmelCase__		  =  1
						lowerCAmelCase__		  =  None
						lowerCAmelCase__		  =  pipe(**a_						).images
						assert image.shape == (1, 64, 64, 3)

						lowerCAmelCase__		  =  image[0, -3:, -3:, -1]

						lowerCAmelCase__		  =  np.array([0.0340, 0.0152, 0.0063, 0.0267, 0.0221, 0.0107, 0.0416, 0.0186, 0.0217]						)

						assert np.abs(image_slice.flatten() - expected_slice						).max() < 2e-2
		@require_torch_a
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  UNetaDModel.from_pretrained('diffusers/consistency_models'    ,subfolder='diffusers_cd_imagenet64_l2'						)
						lowerCAmelCase__		  =  CMStochasticIterativeScheduler(
						    num_train_timesteps=40    ,sigma_min=0.002    ,sigma_max=80.0    ,)
						lowerCAmelCase__		  =  ConsistencyModelPipeline(unet=a_    ,scheduler=a_						)
						pipe.to(torch_device=a_    ,torch_dtype=torch.floataa						)
						pipe.set_progress_bar_config(disable=a_						)

						lowerCAmelCase__		  =  self.get_inputs(get_fixed_latents=a_    ,device=a_						)
						# Ensure usage of flash attention in torch 2.0
						with sdp_kernel(enable_flash=a_    ,enable_math=a_    ,enable_mem_efficient=a_						):
										lowerCAmelCase__		  =  pipe(**a_						).images
						assert image.shape == (1, 64, 64, 3)

						lowerCAmelCase__		  =  image[0, -3:, -3:, -1]

						lowerCAmelCase__		  =  np.array([0.1875, 0.1428, 0.1289, 0.2151, 0.2092, 0.1477, 0.1877, 0.1641, 0.1353]						)

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



		@require_torch_a
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  UNetaDModel.from_pretrained('diffusers/consistency_models'    ,subfolder='diffusers_cd_imagenet64_l2'						)
						lowerCAmelCase__		  =  CMStochasticIterativeScheduler(
						    num_train_timesteps=40    ,sigma_min=0.002    ,sigma_max=80.0    ,)
						lowerCAmelCase__		  =  ConsistencyModelPipeline(unet=a_    ,scheduler=a_						)
						pipe.to(torch_device=a_    ,torch_dtype=torch.floataa						)
						pipe.set_progress_bar_config(disable=a_						)

						lowerCAmelCase__		  =  self.get_inputs(get_fixed_latents=a_    ,device=a_						)
						lowerCAmelCase__		  =  1
						lowerCAmelCase__		  =  None
						# Ensure usage of flash attention in torch 2.0
						with sdp_kernel(enable_flash=a_    ,enable_math=a_    ,enable_mem_efficient=a_						):
										lowerCAmelCase__		  =  pipe(**a_						).images
						assert image.shape == (1, 64, 64, 3)

						lowerCAmelCase__		  =  image[0, -3:, -3:, -1]

						lowerCAmelCase__		  =  np.array([0.1663, 0.1948, 0.2275, 0.1680, 0.1204, 0.1245, 0.1858, 0.1338, 0.2095]						)

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





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 	__snake_case   :
		def __init__(							self    ,a_    ,a_=13    ,a_=7    ,a_=False    ,a_=True    ,a_=False    ,a_=True    ,a_=33    ,a_=32    ,a_=5    ,a_=4    ,a_=37    ,a_="gelu"    ,a_=0.1    ,a_=0.1    ,a_=512    ,a_=16    ,a_=2    ,a_=0.02    ,a_=3    ,a_=4    ,a_=None    ,):







						"""simple docstring"""
						lowerCAmelCase__		  =  parent
						lowerCAmelCase__		  =  batch_size
						lowerCAmelCase__		  =  seq_length
						lowerCAmelCase__		  =  is_training
						lowerCAmelCase__		  =  use_input_mask
						lowerCAmelCase__		  =  use_token_type_ids
						lowerCAmelCase__		  =  use_labels
						lowerCAmelCase__		  =  vocab_size
						lowerCAmelCase__		  =  hidden_size
						lowerCAmelCase__		  =  num_hidden_layers
						lowerCAmelCase__		  =  num_attention_heads
						lowerCAmelCase__		  =  intermediate_size
						lowerCAmelCase__		  =  hidden_act
						lowerCAmelCase__		  =  hidden_dropout_prob
						lowerCAmelCase__		  =  attention_probs_dropout_prob
						lowerCAmelCase__		  =  max_position_embeddings
						lowerCAmelCase__		  =  type_vocab_size
						lowerCAmelCase__		  =  type_sequence_label_size
						lowerCAmelCase__		  =  initializer_range
						lowerCAmelCase__		  =  num_labels
						lowerCAmelCase__		  =  num_choices
						lowerCAmelCase__		  =  scope
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  ids_tensor([self.batch_size, self.seq_length]    ,self.vocab_size						)

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

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

						lowerCAmelCase__		  =  self.get_config()

						return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""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 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  EsmModel(config=a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  model(a_    ,attention_mask=a_						)
						lowerCAmelCase__		  =  model(a_						)
						lowerCAmelCase__		  =  model(a_						)

						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 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  EsmForMaskedLM(config=a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  model(a_    ,attention_mask=a_    ,labels=a_						)
						self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.seq_length, self.vocab_size)						)
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.num_labels
						lowerCAmelCase__		  =  EsmForTokenClassification(config=a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  model(a_    ,attention_mask=a_    ,labels=a_						)
						self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.seq_length, self.num_labels)						)



		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.prepare_config_and_inputs()
						(
						    (
						    lowerCAmelCase__
						)  ,   (
						    lowerCAmelCase__
						)  ,   (
						    lowerCAmelCase__
						)  ,   (
						    lowerCAmelCase__
						)  ,   (
						    lowerCAmelCase__
						)  ,   (
						    lowerCAmelCase__
						)  ,   
						)		  =  config_and_inputs
						lowerCAmelCase__		  =  {'input_ids': input_ids, 'attention_mask': input_mask}
						return config, inputs_dict



@require_torch
class 	__snake_case   ( SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE							, unittest.TestCase ):
		SCREAMING_SNAKE_CASE__								=	False

		SCREAMING_SNAKE_CASE__								=	(
		    (
		        EsmForMaskedLM,
		        EsmModel,
		        EsmForSequenceClassification,
		        EsmForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		SCREAMING_SNAKE_CASE__								=	()
		SCREAMING_SNAKE_CASE__								=	(
		    {
		        'feature-extraction': EsmModel,
		        'fill-mask': EsmForMaskedLM,
		        'text-classification': EsmForSequenceClassification,
		        'token-classification': EsmForTokenClassification,
		        'zero-shot': EsmForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		SCREAMING_SNAKE_CASE__								=	True
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  EsmModelTester(self						)
						lowerCAmelCase__		  =  ConfigTester(self    ,config_class=a_    ,hidden_size=37						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						self.config_tester.run_common_tests()
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
										lowerCAmelCase__		  =  type
										self.model_tester.create_and_check_model(*a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*a_						)
		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCAmelCase__		  =  EsmModel.from_pretrained(a_						)
										self.assertIsNotNone(a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()[0]
						lowerCAmelCase__		  =  EsmEmbeddings(config=a_						)

						lowerCAmelCase__		  =  torch.as_tensor([[12, 31, 13, model.padding_idx]]						)
						lowerCAmelCase__		  =  torch.as_tensor(
						    [
						        [
						            0 + model.padding_idx + 1,
						            1 + model.padding_idx + 1,
						            2 + model.padding_idx + 1,
						            model.padding_idx,
						        ]
						    ]						)
						lowerCAmelCase__		  =  create_position_ids_from_input_ids(a_    ,model.padding_idx						)
						self.assertEqual(position_ids.shape    ,expected_positions.shape						)
						self.assertTrue(torch.all(torch.eq(a_    ,a_						)						)						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()[0]
						lowerCAmelCase__		  =  EsmEmbeddings(config=a_						)

						lowerCAmelCase__		  =  torch.empty(2    ,4    ,30						)
						lowerCAmelCase__		  =  [
						    0 + embeddings.padding_idx + 1,
						    1 + embeddings.padding_idx + 1,
						    2 + embeddings.padding_idx + 1,
						    3 + embeddings.padding_idx + 1,
						]
						lowerCAmelCase__		  =  torch.as_tensor([expected_single_positions, expected_single_positions]						)
						lowerCAmelCase__		  =  embeddings.create_position_ids_from_inputs_embeds(a_						)
						self.assertEqual(position_ids.shape    ,expected_positions.shape						)
						self.assertTrue(torch.all(torch.eq(a_    ,a_						)						)						)
		@unittest.skip('Esm does not support embedding resizing'						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						pass
		@unittest.skip('Esm does not support embedding resizing'						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						pass



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







						"""simple docstring"""
						pass



@require_torch
class 	__snake_case   ( SCREAMING_SNAKE_CASE ):
		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						with torch.no_grad():
										lowerCAmelCase__		  =  EsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'						)
										model.eval()
										lowerCAmelCase__		  =  torch.tensor([[0, 1, 2, 3, 4, 5]]						)
										lowerCAmelCase__		  =  model(a_						)[0]

										lowerCAmelCase__		  =  33

										lowerCAmelCase__		  =  torch.Size((1, 6, vocab_size)						)
										self.assertEqual(output.shape    ,a_						)

										lowerCAmelCase__		  =  torch.tensor(
										    [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]]						)
										self.assertTrue(torch.allclose(output[:, :3, :3]    ,a_    ,atol=1e-4						)						)



		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						with torch.no_grad():
										lowerCAmelCase__		  =  EsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'						)
										model.eval()

										lowerCAmelCase__		  =  torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]						)
										lowerCAmelCase__		  =  model(a_						)[0]
										# compare the actual values for a slice.
										lowerCAmelCase__		  =  torch.tensor(
										    [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]						)
										self.assertTrue(torch.allclose(output[:, :3, :3]    ,a_    ,atol=1e-4						)						)
