

import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionInstructPixaPixPipeline,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class SCREAMING_SNAKE_CASE  (		__snake_case							,						__snake_case							,						__snake_case							,						unittest.TestCase							):






			"""simple docstring"""




			__A						  =  StableDiffusionInstructPixaPixPipeline
			__A						  =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width", "cross_attention_kwargs"}
			__A						  =  TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
			__A						  =  IMAGE_TO_IMAGE_IMAGE_PARAMS
			__A						  =  IMAGE_TO_IMAGE_IMAGE_PARAMS


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				torch.manual_seed(0							)
				snake_case_      		= UNetaDConditionModel(
				    block_out_channels=(32, 64) ,				layers_per_block=2 ,				sample_size=32 ,				in_channels=8 ,				out_channels=4 ,				down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D') ,				up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D') ,				cross_attention_dim=32 ,				)
				snake_case_      		= PNDMScheduler(skip_prk_steps=UpperCamelCase__							)
				torch.manual_seed(0							)
				snake_case_      		= AutoencoderKL(
				    block_out_channels=[32, 64] ,				in_channels=3 ,				out_channels=3 ,				down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D'] ,				up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D'] ,				latent_channels=4 ,				)
				torch.manual_seed(0							)
				snake_case_      		= CLIPTextConfig(
				    bos_token_id=0 ,				eos_token_id=2 ,				hidden_size=32 ,				intermediate_size=37 ,				layer_norm_eps=1E-05 ,				num_attention_heads=4 ,				num_hidden_layers=5 ,				pad_token_id=1 ,				vocab_size=10_00 ,				)
				snake_case_      		= CLIPTextModel(UpperCamelCase__							)
				snake_case_      		= CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'							)

				snake_case_      		= {
				    'unet': unet,
				    'scheduler': scheduler,
				    'vae': vae,
				    'text_encoder': text_encoder,
				    'tokenizer': tokenizer,
				    'safety_checker': None,
				    'feature_extractor': None,
				}
				return components


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase=0							):

				"""simple docstring"""


				snake_case_      		= floats_tensor((1, 3, 32, 32) ,				rng=random.Random(UpperCamelCase__							)							).to(UpperCamelCase__							)
				snake_case_      		= image.cpu().permute(0 ,				2 ,				3 ,				1							)[0]
				snake_case_      		= Image.fromarray(np.uinta(UpperCamelCase__							)							).convert('RGB'							)
				if str(UpperCamelCase__							).startswith('mps'							):
					snake_case_      		= torch.manual_seed(UpperCamelCase__							)
				else:
					snake_case_      		= torch.Generator(device=UpperCamelCase__							).manual_seed(UpperCamelCase__							)
				snake_case_      		= {
				    'prompt': 'A painting of a squirrel eating a burger',
				    'image': image,
				    'generator': generator,
				    'num_inference_steps': 2,
				    'guidance_scale': 6.0,
				    'image_guidance_scale': 1,
				    'output_type': 'numpy',
				}
				return inputs


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= 'cpu'  # ensure determinism for the device-dependent torch.Generator
				snake_case_      		= self.get_dummy_components()
				snake_case_      		= StableDiffusionInstructPixaPixPipeline(**UpperCamelCase__							)
				snake_case_      		= sd_pipe.to(UpperCamelCase__							)
				sd_pipe.set_progress_bar_config(disable=UpperCamelCase__							)

				snake_case_      		= self.get_dummy_inputs(UpperCamelCase__							)
				snake_case_      		= sd_pipe(**UpperCamelCase__							).images
				snake_case_      		= image[0, -3:, -3:, -1]
				assert image.shape == (1, 32, 32, 3)
				snake_case_      		= np.array([0.7526, 0.3750, 0.4547, 0.6117, 0.5866, 0.5016, 0.4327, 0.5642, 0.4815]							)

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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= 'cpu'  # ensure determinism for the device-dependent torch.Generator
				snake_case_      		= self.get_dummy_components()
				snake_case_      		= StableDiffusionInstructPixaPixPipeline(**UpperCamelCase__							)
				snake_case_      		= sd_pipe.to(UpperCamelCase__							)
				sd_pipe.set_progress_bar_config(disable=UpperCamelCase__							)

				snake_case_      		= self.get_dummy_inputs(UpperCamelCase__							)
				snake_case_      		= 'french fries'
				snake_case_      		= sd_pipe(**UpperCamelCase__ ,				negative_prompt=UpperCamelCase__							)
				snake_case_      		= output.images
				snake_case_      		= image[0, -3:, -3:, -1]

				assert image.shape == (1, 32, 32, 3)
				snake_case_      		= np.array([0.7511, 0.3642, 0.4553, 0.6236, 0.5797, 0.5013, 0.4343, 0.5611, 0.4831]							)

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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= 'cpu'  # ensure determinism for the device-dependent torch.Generator
				snake_case_      		= self.get_dummy_components()
				snake_case_      		= StableDiffusionInstructPixaPixPipeline(**UpperCamelCase__							)
				snake_case_      		= sd_pipe.to(UpperCamelCase__							)
				sd_pipe.set_progress_bar_config(disable=UpperCamelCase__							)

				snake_case_      		= self.get_dummy_inputs(UpperCamelCase__							)
				snake_case_      		= [inputs['prompt']] * 2

				snake_case_      		= np.array(inputs['image']							).astype(np.floataa							) / 255.0
				snake_case_      		= torch.from_numpy(UpperCamelCase__							).unsqueeze(0							).to(UpperCamelCase__							)
				snake_case_      		= image / 2 + 0.5
				snake_case_      		= image.permute(0 ,				3 ,				1 ,				2							)
				snake_case_      		= image.repeat(2 ,				1 ,				1 ,				1							)

				snake_case_      		= sd_pipe(**UpperCamelCase__							).images
				snake_case_      		= image[-1, -3:, -3:, -1]

				assert image.shape == (2, 32, 32, 3)
				snake_case_      		= np.array([0.5812, 0.5748, 0.5222, 0.5908, 0.5695, 0.7174, 0.6804, 0.5523, 0.5579]							)

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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= 'cpu'  # ensure determinism for the device-dependent torch.Generator
				snake_case_      		= self.get_dummy_components()
				snake_case_      		= EulerAncestralDiscreteScheduler(
				    beta_start=0.0_0085 ,				beta_end=0.012 ,				beta_schedule='scaled_linear'							)
				snake_case_      		= StableDiffusionInstructPixaPixPipeline(**UpperCamelCase__							)
				snake_case_      		= sd_pipe.to(UpperCamelCase__							)
				sd_pipe.set_progress_bar_config(disable=UpperCamelCase__							)

				snake_case_      		= self.get_dummy_inputs(UpperCamelCase__							)
				snake_case_      		= sd_pipe(**UpperCamelCase__							).images
				snake_case_      		= image[0, -3:, -3:, -1]

				snake_case_      		= [round(UpperCamelCase__ ,				4							) for x in image_slice.flatten().tolist()]
				print(','.join([str(UpperCamelCase__							) for x in slice]							)							)

				assert image.shape == (1, 32, 32, 3)
				snake_case_      		= np.array([0.7417, 0.3842, 0.4732, 0.5776, 0.5891, 0.5139, 0.4052, 0.5673, 0.4986]							)

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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				super().test_inference_batch_single_identical(expected_max_diff=3E-3							)







			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.get_dummy_components()
				snake_case_      		= StableDiffusionInstructPixaPixPipeline(**UpperCamelCase__							)
				snake_case_      		= VaeImageProcessor(do_resize=UpperCamelCase__ ,				do_normalize=UpperCamelCase__							)
				snake_case_      		= pipe.to(UpperCamelCase__							)
				pipe.set_progress_bar_config(disable=UpperCamelCase__							)

				snake_case_      		= pipe(**self.get_dummy_inputs_by_type(UpperCamelCase__ ,				input_image_type='pt'							)							)[0]

				snake_case_      		= components['vae']
				snake_case_      		= self.get_dummy_inputs_by_type(UpperCamelCase__ ,				input_image_type='pt'							)

				for image_param in self.image_latents_params:
					if image_param in inputs.keys():
						snake_case_      		= vae.encode(inputs[image_param]							).latent_dist.mode()

				snake_case_      		= pipe(**UpperCamelCase__							)[0]

				snake_case_      		= np.abs(out - out_latents_inputs							).max()
				self.assertLess(UpperCamelCase__ ,				1E-4 ,				'passing latents as image input generate different result from passing image'							)







@slow
@require_torch_gpu
class SCREAMING_SNAKE_CASE  (		unittest.TestCase							):






			"""simple docstring"""




			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase=0							):

				"""simple docstring"""


				snake_case_      		= torch.manual_seed(UpperCamelCase__							)
				snake_case_      		= load_image(
				    'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg'							)
				snake_case_      		= {
				    'prompt': 'turn him into a cyborg',
				    'image': image,
				    'generator': generator,
				    'num_inference_steps': 3,
				    'guidance_scale': 7.5,
				    'image_guidance_scale': 1.0,
				    'output_type': 'numpy',
				}
				return inputs


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= StableDiffusionInstructPixaPixPipeline.from_pretrained(
				    'timbrooks/instruct-pix2pix' ,				safety_checker=UpperCamelCase__							)
				pipe.to(UpperCamelCase__							)
				pipe.set_progress_bar_config(disable=UpperCamelCase__							)
				pipe.enable_attention_slicing()

				snake_case_      		= self.get_inputs()
				snake_case_      		= pipe(**UpperCamelCase__							).images
				snake_case_      		= image[0, -3:, -3:, -1].flatten()

				assert image.shape == (1, 5_12, 5_12, 3)
				snake_case_      		= np.array([0.5902, 0.6015, 0.6027, 0.5983, 0.6092, 0.6061, 0.5765, 0.5785, 0.5555]							)

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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= StableDiffusionInstructPixaPixPipeline.from_pretrained(
				    'timbrooks/instruct-pix2pix' ,				safety_checker=UpperCamelCase__							)
				snake_case_      		= LMSDiscreteScheduler.from_config(pipe.scheduler.config							)
				pipe.to(UpperCamelCase__							)
				pipe.set_progress_bar_config(disable=UpperCamelCase__							)
				pipe.enable_attention_slicing()

				snake_case_      		= self.get_inputs()
				snake_case_      		= pipe(**UpperCamelCase__							).images
				snake_case_      		= image[0, -3:, -3:, -1].flatten()

				assert image.shape == (1, 5_12, 5_12, 3)
				snake_case_      		= np.array([0.6578, 0.6817, 0.6972, 0.6761, 0.6856, 0.6916, 0.6428, 0.6516, 0.6301]							)

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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= StableDiffusionInstructPixaPixPipeline.from_pretrained(
				    'timbrooks/instruct-pix2pix' ,				safety_checker=UpperCamelCase__							)
				snake_case_      		= DDIMScheduler.from_config(pipe.scheduler.config							)
				pipe.to(UpperCamelCase__							)
				pipe.set_progress_bar_config(disable=UpperCamelCase__							)
				pipe.enable_attention_slicing()

				snake_case_      		= self.get_inputs()
				snake_case_      		= pipe(**UpperCamelCase__							).images
				snake_case_      		= image[0, -3:, -3:, -1].flatten()

				assert image.shape == (1, 5_12, 5_12, 3)
				snake_case_      		= np.array([0.3828, 0.3834, 0.3818, 0.3792, 0.3865, 0.3752, 0.3792, 0.3847, 0.3753]							)

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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= 0

				def callback_fn(__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							) -> None:
					snake_case_      		= True
					nonlocal number_of_steps
					number_of_steps += 1
					if step == 1:
						snake_case_      		= latents.detach().cpu().numpy()
						assert latents.shape == (1, 4, 64, 64)
						snake_case_      		= latents[0, -3:, -3:, -1]
						snake_case_      		= np.array([-0.2463, -0.4644, -0.9756, 1.5176, 1.4414, 0.7866, 0.9897, 0.8521, 0.7983]							)

						assert np.abs(latents_slice.flatten() - expected_slice							).max() < 5E-2
					elif step == 2:
						snake_case_      		= latents.detach().cpu().numpy()
						assert latents.shape == (1, 4, 64, 64)
						snake_case_      		= latents[0, -3:, -3:, -1]
						snake_case_      		= np.array([-0.2644, -0.4626, -0.9653, 1.5176, 1.4551, 0.7686, 0.9805, 0.8452, 0.8115]							)

						assert np.abs(latents_slice.flatten() - expected_slice							).max() < 5E-2

				snake_case_      		= False

				snake_case_      		= StableDiffusionInstructPixaPixPipeline.from_pretrained(
				    'timbrooks/instruct-pix2pix' ,				safety_checker=UpperCamelCase__ ,				torch_dtype=torch.floataa							)
				snake_case_      		= pipe.to(UpperCamelCase__							)
				pipe.set_progress_bar_config(disable=UpperCamelCase__							)
				pipe.enable_attention_slicing()

				snake_case_      		= self.get_inputs()
				pipe(**UpperCamelCase__ ,				callback=UpperCamelCase__ ,				callback_steps=1							)
				assert callback_fn.has_been_called
				assert number_of_steps == 3


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				torch.cuda.empty_cache()
				torch.cuda.reset_max_memory_allocated()
				torch.cuda.reset_peak_memory_stats()

				snake_case_      		= StableDiffusionInstructPixaPixPipeline.from_pretrained(
				    'timbrooks/instruct-pix2pix' ,				safety_checker=UpperCamelCase__ ,				torch_dtype=torch.floataa							)
				snake_case_      		= pipe.to(UpperCamelCase__							)
				pipe.set_progress_bar_config(disable=UpperCamelCase__							)
				pipe.enable_attention_slicing(1							)
				pipe.enable_sequential_cpu_offload()

				snake_case_      		= self.get_inputs()
				snake_case_      		= pipe(**UpperCamelCase__							)

				snake_case_      		= torch.cuda.max_memory_allocated()
				# make sure that less than 2.2 GB is allocated
				assert mem_bytes < 2.2 * 10**9







			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.get_inputs()
				# resize to resolution that is divisible by 8 but not 16 or 32
				snake_case_      		= inputs['image'].resize((5_04, 5_04)							)

				snake_case_      		= 'timbrooks/instruct-pix2pix'
				snake_case_      		= StableDiffusionInstructPixaPixPipeline.from_pretrained(
				    UpperCamelCase__ ,				safety_checker=UpperCamelCase__ ,				)
				pipe.to(UpperCamelCase__							)
				pipe.set_progress_bar_config(disable=UpperCamelCase__							)
				pipe.enable_attention_slicing()

				snake_case_      		= pipe(**UpperCamelCase__							)
				snake_case_      		= output.images[0]

				snake_case_      		= image[2_55:2_58, 3_83:3_86, -1]

				assert image.shape == (5_04, 5_04, 3)
				snake_case_      		= np.array([0.2726, 0.2529, 0.2664, 0.2655, 0.2641, 0.2642, 0.2591, 0.2649, 0.2590]							)

				assert np.abs(image_slice.flatten() - expected_slice							).max() < 5E-3








import unittest

from transformers import MraConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_torch_available():
		import torch

		from transformers import (
		    MraForMaskedLM,
		    MraForMultipleChoice,
		    MraForQuestionAnswering,
		    MraForSequenceClassification,
		    MraForTokenClassification,
		    MraModel,
		)
		from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST


class SCREAMING_SNAKE_CASE  :






			"""simple docstring"""




			def __init__(    self ,				__UpperCamelCase ,				__UpperCamelCase=2 ,				__UpperCamelCase=8 ,				__UpperCamelCase=True ,				__UpperCamelCase=True ,				__UpperCamelCase=True ,				__UpperCamelCase=True ,				__UpperCamelCase=99 ,				__UpperCamelCase=16 ,				__UpperCamelCase=5 ,				__UpperCamelCase=2 ,				__UpperCamelCase=36 ,				__UpperCamelCase="gelu" ,				__UpperCamelCase=0.0 ,				__UpperCamelCase=0.0 ,				__UpperCamelCase=5_12 ,				__UpperCamelCase=16 ,				__UpperCamelCase=2 ,				__UpperCamelCase=0.02 ,				__UpperCamelCase=3 ,				__UpperCamelCase=4 ,				__UpperCamelCase=None ,				):

				"""simple docstring"""


				snake_case_      		= parent
				snake_case_      		= batch_size
				snake_case_      		= seq_length
				snake_case_      		= is_training
				snake_case_      		= use_input_mask
				snake_case_      		= use_token_type_ids
				snake_case_      		= use_labels
				snake_case_      		= vocab_size
				snake_case_      		= hidden_size
				snake_case_      		= num_hidden_layers
				snake_case_      		= num_attention_heads
				snake_case_      		= intermediate_size
				snake_case_      		= hidden_act
				snake_case_      		= hidden_dropout_prob
				snake_case_      		= attention_probs_dropout_prob
				snake_case_      		= max_position_embeddings
				snake_case_      		= type_vocab_size
				snake_case_      		= type_sequence_label_size
				snake_case_      		= initializer_range
				snake_case_      		= num_labels
				snake_case_      		= num_choices
				snake_case_      		= scope


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= ids_tensor([self.batch_size, self.seq_length] ,				self.vocab_size							)

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

				snake_case_      		= None
				if self.use_token_type_ids:
					snake_case_      		= ids_tensor([self.batch_size, self.seq_length] ,				self.type_vocab_size							)

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

				snake_case_      		= self.get_config()

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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				return MraConfig(
				    vocab_size=self.vocab_size ,				hidden_size=self.hidden_size ,				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 ,				is_decoder=__UpperCamelCase ,				initializer_range=self.initializer_range ,				)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.get_config()
				snake_case_      		= 3_00
				return config


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				(
				    (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       
				)      		= self.prepare_config_and_inputs()

				snake_case_      		= True
				snake_case_      		= floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
				snake_case_      		= ids_tensor([self.batch_size, self.seq_length] ,				vocab_size=2							)

				return (
				    config,
				    input_ids,
				    token_type_ids,
				    input_mask,
				    sequence_labels,
				    token_labels,
				    choice_labels,
				    encoder_hidden_states,
				    encoder_attention_mask,
				)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= MraModel(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase							)
				snake_case_      		= model(__UpperCamelCase ,				token_type_ids=__UpperCamelCase							)
				snake_case_      		= model(__UpperCamelCase							)
				self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, self.seq_length, self.hidden_size)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				):

				"""simple docstring"""


				snake_case_      		= True
				snake_case_      		= MraModel(__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(
				    __UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				encoder_hidden_states=__UpperCamelCase ,				encoder_attention_mask=__UpperCamelCase ,				)
				snake_case_      		= model(
				    __UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				encoder_hidden_states=__UpperCamelCase ,				)
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase							)
				self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, self.seq_length, self.hidden_size)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= MraForMaskedLM(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__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 ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= MraForQuestionAnswering(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(
				    __UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				start_positions=__UpperCamelCase ,				end_positions=__UpperCamelCase ,				)
				self.parent.assertEqual(result.start_logits.shape ,				(self.batch_size, self.seq_length)							)
				self.parent.assertEqual(result.end_logits.shape ,				(self.batch_size, self.seq_length)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= self.num_labels
				snake_case_      		= MraForSequenceClassification(__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				labels=__UpperCamelCase							)
				self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_labels)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= self.num_labels
				snake_case_      		= MraForTokenClassification(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				labels=__UpperCamelCase							)
				self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.num_labels)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= self.num_choices
				snake_case_      		= MraForMultipleChoice(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= input_ids.unsqueeze(1							).expand(-1 ,				self.num_choices ,				-1							).contiguous()
				snake_case_      		= token_type_ids.unsqueeze(1							).expand(-1 ,				self.num_choices ,				-1							).contiguous()
				snake_case_      		= input_mask.unsqueeze(1							).expand(-1 ,				self.num_choices ,				-1							).contiguous()
				snake_case_      		= model(
				    __UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				labels=__UpperCamelCase ,				)
				self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_choices)							)







			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.prepare_config_and_inputs()
				(
				    (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       
				)      		= config_and_inputs
				snake_case_      		= {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
				return config, inputs_dict







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






			"""simple docstring"""




			__A						  =  (
			    (
			        MraModel,
			        MraForMaskedLM,
			        MraForMultipleChoice,
			        MraForQuestionAnswering,
			        MraForSequenceClassification,
			        MraForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			__A						  =  False
			__A						  =  False
			__A						  =  False
			__A						  =  False

			__A						  =  ()


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= MraModelTester(self							)
				snake_case_      		= ConfigTester(self ,				config_class=__UpperCamelCase ,				hidden_size=37							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				self.config_tester.run_common_tests()


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				for type in ["absolute", "relative_key", "relative_key_query"]:
					snake_case_      		= type
					self.model_tester.create_and_check_model(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_lm(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_multiple_choice(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_question_answering(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_sequence_classification(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*__UpperCamelCase							)


			@slow
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					snake_case_      		= MraModel.from_pretrained(__UpperCamelCase							)
					self.assertIsNotNone(__UpperCamelCase							)







			@unittest.skip(reason='MRA does not output attentions'							)
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				return







@require_torch
class SCREAMING_SNAKE_CASE  (		unittest.TestCase							):






			"""simple docstring"""




			@slow
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= MraModel.from_pretrained('uw-madison/mra-base-512-4'							)
				snake_case_      		= torch.arange(2_56							).unsqueeze(0							)

				with torch.no_grad():
					snake_case_      		= model(__UpperCamelCase							)[0]

				snake_case_      		= torch.Size((1, 2_56, 7_68)							)
				self.assertEqual(output.shape ,				__UpperCamelCase							)

				snake_case_      		= torch.tensor(
				    [[[-0.0140, 0.0830, -0.0381], [0.1546, 0.1402, 0.0220], [0.1162, 0.0851, 0.0165]]]							)

				self.assertTrue(torch.allclose(output[:, :3, :3] ,				__UpperCamelCase ,				atol=1E-4							)							)


			@slow
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= MraForMaskedLM.from_pretrained('uw-madison/mra-base-512-4'							)
				snake_case_      		= torch.arange(2_56							).unsqueeze(0							)

				with torch.no_grad():
					snake_case_      		= model(__UpperCamelCase							)[0]

				snake_case_      		= 5_02_65

				snake_case_      		= torch.Size((1, 2_56, vocab_size)							)
				self.assertEqual(output.shape ,				__UpperCamelCase							)

				snake_case_      		= torch.tensor(
				    [[[9.2595, -3.6038, 11.8819], [9.3869, -3.2693, 11.0956], [11.8524, -3.4938, 13.1210]]]							)

				self.assertTrue(torch.allclose(output[:, :3, :3] ,				__UpperCamelCase ,				atol=1E-4							)							)







			@slow
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= MraForMaskedLM.from_pretrained('uw-madison/mra-base-4096-8-d3'							)
				snake_case_      		= torch.arange(40_96							).unsqueeze(0							)

				with torch.no_grad():
					snake_case_      		= model(__UpperCamelCase							)[0]

				snake_case_      		= 5_02_65

				snake_case_      		= torch.Size((1, 40_96, vocab_size)							)
				self.assertEqual(output.shape ,				__UpperCamelCase							)

				snake_case_      		= torch.tensor(
				    [[[5.4789, -2.3564, 7.5064], [7.9067, -1.3369, 9.9668], [9.0712, -1.8106, 7.0380]]]							)

				self.assertTrue(torch.allclose(output[:, :3, :3] ,				__UpperCamelCase ,				atol=1E-4							)							)
