'''simple docstring'''




import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyImgaImgPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()





class     __magic_name__						(    _UpperCAmelCase,     unittest.TestCase):
			UpperCamelCase__				=       KandinskyImgaImgPipeline
			UpperCamelCase__				=       ['''prompt''', '''image_embeds''', '''negative_image_embeds''', '''image''']
			UpperCamelCase__				=       [
			    '''prompt''',
			    '''negative_prompt''',
			    '''image_embeds''',
			    '''negative_image_embeds''',
			    '''image''',
			]
			UpperCamelCase__				=       [
			    '''generator''',
			    '''height''',
			    '''width''',
			    '''strength''',
			    '''guidance_scale''',
			    '''negative_prompt''',
			    '''num_inference_steps''',
			    '''return_dict''',
			    '''guidance_scale''',
			    '''num_images_per_prompt''',
			    '''output_type''',
			    '''return_dict''',
			]
			UpperCamelCase__				=       False







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict    ):
										return 32







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  str    ):
										return 32







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]    ):
										return self.time_input_dim







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  int    ):
										return self.time_input_dim * 4







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict    ):
										return 100







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict    ):
										lowercase_					:			str										=       XLMRobertaTokenizerFast.from_pretrained("""YiYiXu/tiny-random-mclip-base"""    )
										return tokenizer







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[str]    ):
										torch.manual_seed(0    )
										lowercase_					:			List[str]										=       MCLIPConfig(
										    numDims=self.cross_attention_dim				,			transformerDimensions=self.text_embedder_hidden_size				,			hidden_size=self.text_embedder_hidden_size				,			intermediate_size=37				,			num_attention_heads=4				,			num_hidden_layers=5				,			vocab_size=1005				,			)

										lowercase_					:			int										=       MultilingualCLIP(lowercase_    )
										lowercase_					:			Union[str, Any]										=       text_encoder.eval()

										return text_encoder







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict    ):
										torch.manual_seed(0    )

										lowercase_					:			Union[str, Any]										=       {
										    """in_channels""": 4,
										    # Out channels is double in channels because predicts mean and variance
										    """out_channels""": 8,
										    """addition_embed_type""": """text_image""",
										    """down_block_types""": ("""ResnetDownsampleBlock2D""", """SimpleCrossAttnDownBlock2D"""),
										    """up_block_types""": ("""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""),
										    """mid_block_type""": """UNetMidBlock2DSimpleCrossAttn""",
										    """block_out_channels""": (self.block_out_channels_a, self.block_out_channels_a * 2),
										    """layers_per_block""": 1,
										    """encoder_hid_dim""": self.text_embedder_hidden_size,
										    """encoder_hid_dim_type""": """text_image_proj""",
										    """cross_attention_dim""": self.cross_attention_dim,
										    """attention_head_dim""": 4,
										    """resnet_time_scale_shift""": """scale_shift""",
										    """class_embed_type""": None,
										}

										lowercase_					:			List[str]										=       UNetaDConditionModel(**lowercase_    )
										return model







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[str]    ):
										return {
										    "block_out_channels": [32, 64],
										    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
										    "in_channels": 3,
										    "latent_channels": 4,
										    "layers_per_block": 1,
										    "norm_num_groups": 8,
										    "norm_type": "spatial",
										    "num_vq_embeddings": 12,
										    "out_channels": 3,
										    "up_block_types": [
										        "AttnUpDecoderBlock2D",
										        "UpDecoderBlock2D",
										    ],
										    "vq_embed_dim": 4,
										}







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict    ):
										torch.manual_seed(0    )
										lowercase_					:			int										=       VQModel(**self.dummy_movq_kwargs    )
										return model







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										lowercase_					:			Dict										=       self.dummy_text_encoder
										lowercase_					:			Tuple										=       self.dummy_tokenizer
										lowercase_					:			str										=       self.dummy_unet
										lowercase_					:			Tuple										=       self.dummy_movq

										lowercase_					:			List[Any]										=       {
										    """num_train_timesteps""": 1000,
										    """beta_schedule""": """linear""",
										    """beta_start""": 0.0_00_85,
										    """beta_end""": 0.0_12,
										    """clip_sample""": False,
										    """set_alpha_to_one""": False,
										    """steps_offset""": 0,
										    """prediction_type""": """epsilon""",
										    """thresholding""": False,
										}

										lowercase_					:			Tuple										=       DDIMScheduler(**lowercase_    )

										lowercase_					:			List[str]										=       {
										    """text_encoder""": text_encoder,
										    """tokenizer""": tokenizer,
										    """unet""": unet,
										    """scheduler""": scheduler,
										    """movq""": movq,
										}

										return components







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple				,			lowercase_				:  Any				,			lowercase_				:  List[str]=0    ):
										lowercase_					:			List[Any]										=       floats_tensor((1, self.cross_attention_dim)				,			rng=random.Random(lowercase_    )    ).to(lowercase_    )
										lowercase_					:			Any										=       floats_tensor((1, self.cross_attention_dim)				,			rng=random.Random(seed + 1    )    ).to(lowercase_    )
										# create init_image
										lowercase_					:			int										=       floats_tensor((1, 3, 64, 64)				,			rng=random.Random(lowercase_    )    ).to(lowercase_    )
										lowercase_					:			List[str]										=       image.cpu().permute(0				,			2				,			3				,			1    )[0]
										lowercase_					:			Any										=       Image.fromarray(np.uinta(lowercase_    )    ).convert("""RGB"""    ).resize((256, 256)    )

										if str(lowercase_    ).startswith("""mps"""    ):
																	lowercase_					:			Union[str, Any]										=       torch.manual_seed(lowercase_    )
										else:
																	lowercase_					:			Optional[int]										=       torch.Generator(device=lowercase_    ).manual_seed(lowercase_    )
										lowercase_					:			Dict										=       {
										    """prompt""": """horse""",
										    """image""": init_image,
										    """image_embeds""": image_embeds,
										    """negative_image_embeds""": negative_image_embeds,
										    """generator""": generator,
										    """height""": 64,
										    """width""": 64,
										    """num_inference_steps""": 10,
										    """guidance_scale""": 7.0,
										    """strength""": 0.2,
										    """output_type""": """np""",
										}
										return inputs







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										lowercase_					:			Optional[Any]										=       """cpu"""

										lowercase_					:			List[str]										=       self.get_dummy_components()

										lowercase_					:			Optional[Any]										=       self.pipeline_class(**lowercase_    )
										lowercase_					:			List[Any]										=       pipe.to(lowercase_    )

										pipe.set_progress_bar_config(disable=lowercase_    )

										lowercase_					:			Tuple										=       pipe(**self.get_dummy_inputs(lowercase_    )    )
										lowercase_					:			Optional[int]										=       output.images

										lowercase_					:			List[str]										=       pipe(
										    **self.get_dummy_inputs(lowercase_    )				,			return_dict=lowercase_				,			)[0]

										lowercase_					:			Tuple										=       image[0, -3:, -3:, -1]
										lowercase_					:			Optional[Any]										=       image_from_tuple[0, -3:, -3:, -1]

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

										lowercase_					:			Dict										=       np.array(
										    [0.61_47_49_43, 0.6_07_35_39, 0.43_30_85_44, 0.5_92_82_69, 0.47_49_35_95, 0.46_75_59_73, 0.4_61_38_38, 0.45_36_87_97, 0.50_11_92_33]    )
										assert (
										    np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2
										), f''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''
										assert (
										    np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() < 1E-2
										), f''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''





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







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple    ):
										lowercase_					:			Tuple										=       load_numpy(
										    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
										    """/kandinsky/kandinsky_img2img_frog.npy"""    )

										lowercase_					:			Optional[Any]										=       load_image(
										    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/kandinsky/cat.png"""    )
										lowercase_					:			int										=       """A red cartoon frog, 4k"""

										lowercase_					:			List[Any]										=       KandinskyPriorPipeline.from_pretrained(
										    """kandinsky-community/kandinsky-2-1-prior"""				,			torch_dtype=torch.floataa    )
										pipe_prior.to(lowercase_    )

										lowercase_					:			Optional[Any]										=       KandinskyImgaImgPipeline.from_pretrained(
										    """kandinsky-community/kandinsky-2-1"""				,			torch_dtype=torch.floataa    )
										lowercase_					:			Optional[Any]										=       pipeline.to(lowercase_    )

										pipeline.set_progress_bar_config(disable=lowercase_    )

										lowercase_					:			str										=       torch.Generator(device="""cpu"""    ).manual_seed(0    )
										lowercase_						,						lowercase_					:			str										=       pipe_prior(
										    lowercase_				,			generator=lowercase_				,			num_inference_steps=5				,			negative_prompt=""""""				,			).to_tuple()

										lowercase_					:			Optional[int]										=       pipeline(
										    lowercase_				,			image=lowercase_				,			image_embeds=lowercase_				,			negative_image_embeds=lowercase_				,			generator=lowercase_				,			num_inference_steps=100				,			height=768				,			width=768				,			strength=0.2				,			output_type="""np"""				,			)

										lowercase_					:			Dict										=       output.images[0]

										assert image.shape == (768, 768, 3)

										assert_mean_pixel_difference(lowercase_				,			lowercase_    )






'''simple docstring'''




import unittest

from transformers import AutoTokenizer, NystromformerConfig, 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, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import (
						    NystromformerForMaskedLM,
						    NystromformerForMultipleChoice,
						    NystromformerForQuestionAnswering,
						    NystromformerForSequenceClassification,
						    NystromformerForTokenClassification,
						    NystromformerModel,
						)
						from transformers.models.nystromformer.modeling_nystromformer import NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST





class     __magic_name__						:







			def __init__(	self				:  int				,			lowercase_				:  Optional[int]				,			lowercase_				:  Any=13				,			lowercase_				:  List[str]=7				,			lowercase_				:  List[Any]=True				,			lowercase_				:  str=True				,			lowercase_				:  Dict=True				,			lowercase_				:  List[str]=True				,			lowercase_				:  List[str]=99				,			lowercase_				:  Dict=32				,			lowercase_				:  List[Any]=5				,			lowercase_				:  List[str]=4				,			lowercase_				:  Dict=37				,			lowercase_				:  List[Any]="gelu"				,			lowercase_				:  Dict=0.1				,			lowercase_				:  Any=0.1				,			lowercase_				:  int=512				,			lowercase_				:  Tuple=16				,			lowercase_				:  str=2				,			lowercase_				:  Optional[Any]=0.02				,			lowercase_				:  Any=3				,			lowercase_				:  Any=4				,			lowercase_				:  Dict=None				,			):
										lowercase_					:			Tuple										=       parent
										lowercase_					:			Tuple										=       batch_size
										lowercase_					:			Optional[int]										=       seq_length
										lowercase_					:			Union[str, Any]										=       is_training
										lowercase_					:			int										=       use_input_mask
										lowercase_					:			Union[str, Any]										=       use_token_type_ids
										lowercase_					:			Tuple										=       use_labels
										lowercase_					:			Tuple										=       vocab_size
										lowercase_					:			int										=       hidden_size
										lowercase_					:			int										=       num_hidden_layers
										lowercase_					:			Optional[int]										=       num_attention_heads
										lowercase_					:			Union[str, Any]										=       intermediate_size
										lowercase_					:			List[Any]										=       hidden_act
										lowercase_					:			int										=       hidden_dropout_prob
										lowercase_					:			Union[str, Any]										=       attention_probs_dropout_prob
										lowercase_					:			List[Any]										=       max_position_embeddings
										lowercase_					:			Union[str, Any]										=       type_vocab_size
										lowercase_					:			List[Any]										=       type_sequence_label_size
										lowercase_					:			Optional[int]										=       initializer_range
										lowercase_					:			str										=       num_labels
										lowercase_					:			int										=       num_choices
										lowercase_					:			List[Any]										=       scope







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple    ):
										lowercase_					:			int										=       ids_tensor([self.batch_size, self.seq_length]				,			self.vocab_size    )

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

										lowercase_					:			Optional[int]										=       None
										if self.use_token_type_ids:
																	lowercase_					:			Tuple										=       ids_tensor([self.batch_size, self.seq_length]				,			self.type_vocab_size    )

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

										lowercase_					:			Tuple										=       self.get_config()

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







			def 	SCREAMING_SNAKE_CASE_     (	self				:  int    ):
										return NystromformerConfig(
										    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=lowercase_				,			initializer_range=self.initializer_range				,			)







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]				,			lowercase_				:  Any				,			lowercase_				:  List[str]				,			lowercase_				:  Union[str, Any]				,			lowercase_				:  Union[str, Any]				,			lowercase_				:  Tuple				,			lowercase_				:  Tuple				,			lowercase_				:  Optional[int]    ):
										lowercase_					:			Optional[Any]										=       NystromformerModel(config=lowercase_    )
										model.to(lowercase_    )
										model.eval()
										lowercase_					:			Optional[int]										=       model(lowercase_				,			attention_mask=lowercase_				,			token_type_ids=lowercase_    )
										lowercase_					:			Optional[Any]										=       model(lowercase_				,			token_type_ids=lowercase_    )
										lowercase_					:			Union[str, Any]										=       model(lowercase_    )
										self.parent.assertEqual(result.last_hidden_state.shape				,			(self.batch_size, self.seq_length, self.hidden_size)    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  int				,			lowercase_				:  Tuple				,			lowercase_				:  List[Any]				,			lowercase_				:  Union[str, Any]				,			lowercase_				:  List[str]				,			lowercase_				:  List[Any]				,			lowercase_				:  Optional[int]				,			lowercase_				:  Any    ):
										lowercase_					:			List[Any]										=       NystromformerForMaskedLM(config=lowercase_    )
										model.to(lowercase_    )
										model.eval()
										lowercase_					:			Optional[int]										=       model(lowercase_				,			attention_mask=lowercase_				,			token_type_ids=lowercase_				,			labels=lowercase_    )
										self.parent.assertEqual(result.logits.shape				,			(self.batch_size, self.seq_length, self.vocab_size)    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple				,			lowercase_				:  int				,			lowercase_				:  List[Any]				,			lowercase_				:  Any				,			lowercase_				:  str				,			lowercase_				:  Dict				,			lowercase_				:  Any				,			lowercase_				:  Tuple    ):
										lowercase_					:			Any										=       NystromformerForQuestionAnswering(config=lowercase_    )
										model.to(lowercase_    )
										model.eval()
										lowercase_					:			Union[str, Any]										=       model(
										    lowercase_				,			attention_mask=lowercase_				,			token_type_ids=lowercase_				,			start_positions=lowercase_				,			end_positions=lowercase_				,			)
										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 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple				,			lowercase_				:  List[str]				,			lowercase_				:  List[Any]				,			lowercase_				:  Union[str, Any]				,			lowercase_				:  str				,			lowercase_				:  Optional[Any]				,			lowercase_				:  str				,			lowercase_				:  int    ):
										lowercase_					:			Any										=       self.num_labels
										lowercase_					:			Union[str, Any]										=       NystromformerForSequenceClassification(lowercase_    )
										model.to(lowercase_    )
										model.eval()
										lowercase_					:			Any										=       model(lowercase_				,			attention_mask=lowercase_				,			token_type_ids=lowercase_				,			labels=lowercase_    )
										self.parent.assertEqual(result.logits.shape				,			(self.batch_size, self.num_labels)    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]				,			lowercase_				:  List[str]				,			lowercase_				:  int				,			lowercase_				:  List[str]				,			lowercase_				:  Dict				,			lowercase_				:  Tuple				,			lowercase_				:  Union[str, Any]				,			lowercase_				:  List[str]    ):
										lowercase_					:			int										=       self.num_labels
										lowercase_					:			int										=       NystromformerForTokenClassification(config=lowercase_    )
										model.to(lowercase_    )
										model.eval()
										lowercase_					:			Tuple										=       model(lowercase_				,			attention_mask=lowercase_				,			token_type_ids=lowercase_				,			labels=lowercase_    )
										self.parent.assertEqual(result.logits.shape				,			(self.batch_size, self.seq_length, self.num_labels)    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]				,			lowercase_				:  Tuple				,			lowercase_				:  Dict				,			lowercase_				:  int				,			lowercase_				:  List[str]				,			lowercase_				:  Optional[Any]				,			lowercase_				:  str				,			lowercase_				:  Union[str, Any]    ):
										lowercase_					:			str										=       self.num_choices
										lowercase_					:			Union[str, Any]										=       NystromformerForMultipleChoice(config=lowercase_    )
										model.to(lowercase_    )
										model.eval()
										lowercase_					:			str										=       input_ids.unsqueeze(1    ).expand(-1				,			self.num_choices				,			-1    ).contiguous()
										lowercase_					:			Tuple										=       token_type_ids.unsqueeze(1    ).expand(-1				,			self.num_choices				,			-1    ).contiguous()
										lowercase_					:			Dict										=       input_mask.unsqueeze(1    ).expand(-1				,			self.num_choices				,			-1    ).contiguous()
										lowercase_					:			Union[str, Any]										=       model(
										    lowercase_				,			attention_mask=lowercase_				,			token_type_ids=lowercase_				,			labels=lowercase_				,			)
										self.parent.assertEqual(result.logits.shape				,			(self.batch_size, self.num_choices)    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict    ):
										lowercase_					:			Union[str, Any]										=       self.prepare_config_and_inputs()
										(
										    (
										    lowercase_
										)						,						(
										    lowercase_
										)						,						(
										    lowercase_
										)						,						(
										    lowercase_
										)						,						(
										    lowercase_
										)						,						(
										    lowercase_
										)						,						(
										    lowercase_
										)						,						
										)					:			Tuple										=       config_and_inputs
										lowercase_					:			Optional[Any]										=       {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
										return config, inputs_dict





@require_torch
class     __magic_name__						(    _UpperCAmelCase,     _UpperCAmelCase,     unittest.TestCase):
			UpperCamelCase__				=       (
			    (
			        NystromformerModel,
			        NystromformerForMaskedLM,
			        NystromformerForMultipleChoice,
			        NystromformerForQuestionAnswering,
			        NystromformerForSequenceClassification,
			        NystromformerForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			UpperCamelCase__				=       (
			    {
			        '''feature-extraction''': NystromformerModel,
			        '''fill-mask''': NystromformerForMaskedLM,
			        '''question-answering''': NystromformerForQuestionAnswering,
			        '''text-classification''': NystromformerForSequenceClassification,
			        '''token-classification''': NystromformerForTokenClassification,
			        '''zero-shot''': NystromformerForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			UpperCamelCase__				=       False
			UpperCamelCase__				=       False







			def 	SCREAMING_SNAKE_CASE_     (	self				:  int    ):
										lowercase_					:			Any										=       NystromformerModelTester(self    )
										lowercase_					:			Optional[Any]										=       ConfigTester(self				,			config_class=lowercase_				,			hidden_size=37    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple    ):
										self.config_tester.run_common_tests()







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict    ):
										lowercase_					:			Dict										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]    ):
										lowercase_					:			Union[str, Any]										=       self.model_tester.prepare_config_and_inputs()
										for type in ["absolute", "relative_key", "relative_key_query"]:
																	lowercase_					:			int										=       type
																	self.model_tester.create_and_check_model(*lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]    ):
										lowercase_					:			Tuple										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										lowercase_					:			int										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_multiple_choice(*lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]    ):
										lowercase_					:			Union[str, Any]										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_question_answering(*lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]    ):
										lowercase_					:			Any										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_sequence_classification(*lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]    ):
										lowercase_					:			Optional[int]										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*lowercase_    )







			@slow
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	lowercase_					:			List[Any]										=       NystromformerModel.from_pretrained(lowercase_    )
																	self.assertIsNotNone(lowercase_    )





@require_torch
class     __magic_name__						(    unittest.TestCase):







			@slow
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]    ):
										lowercase_					:			List[str]										=       NystromformerModel.from_pretrained("""uw-madison/nystromformer-512"""    )
										lowercase_					:			List[str]										=       torch.tensor([[0, 1, 2, 3, 4, 5]]    )

										with torch.no_grad():
																	lowercase_					:			Tuple										=       model(lowercase_    )[0]

										lowercase_					:			Tuple										=       torch.Size((1, 6, 768)    )
										self.assertEqual(output.shape				,			lowercase_    )

										lowercase_					:			Dict										=       torch.tensor(
										    [[[-0.45_32, -0.09_36, 0.51_37], [-0.26_76, 0.06_28, 0.61_86], [-0.36_29, -0.17_26, 0.47_16]]]    )

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







			@slow
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										lowercase_					:			Optional[int]										=       """the [MASK] of Belgium is Brussels"""

										lowercase_					:			Optional[Any]										=       AutoTokenizer.from_pretrained("""uw-madison/nystromformer-512"""    )
										lowercase_					:			List[Any]										=       NystromformerForMaskedLM.from_pretrained("""uw-madison/nystromformer-512"""    )

										lowercase_					:			str										=       tokenizer(lowercase_				,			return_tensors="""pt"""    )

										with torch.no_grad():
																	lowercase_					:			Tuple										=       model(encoding.input_ids    ).logits

										lowercase_					:			Optional[int]										=       token_logits[:, 2, :].argmax(-1    )[0]

										self.assertEqual(tokenizer.decode(lowercase_    )				,			"""capital"""    )





