




'''simple docstring'''
import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()






class 			SCREAMING_SNAKE_CASE(							unittest.TestCase	):

       def        snake_case__     (       self	)       ->   Optional[Any]:

           """simple docstring"""



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

       @property
       def        snake_case__     (       self	)       ->   str:

           """simple docstring"""



           __lowercase     							=		1
           __lowercase     							=		3
           __lowercase     							=		(32, 32)

           __lowercase     							=		floats_tensor((batch_size, num_channels) + sizes   ,						rng=random.Random(0	)	).to(SCREAMING_SNAKE_CASE_	)
           return image

       @property
       def        snake_case__     (       self	)       ->   Optional[int]:

           """simple docstring"""



           torch.manual_seed(0	)
           __lowercase     							=		UNetaDConditionModel(
               block_out_channels=(32, 64)   ,						layers_per_block=2   ,						sample_size=32   ,						in_channels=4   ,						out_channels=4   ,						down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")   ,						up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")   ,						cross_attention_dim=32   ,						)
           return model

       @property
       def        snake_case__     (       self	)       ->   Optional[int]:

           """simple docstring"""



           torch.manual_seed(0	)
           __lowercase     							=		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   ,						)
           return model

       @property
       def        snake_case__     (       self	)       ->   int:

           """simple docstring"""



           torch.manual_seed(0	)
           __lowercase     							=		CLIPTextConfig(
               bos_token_id=0   ,						eos_token_id=2   ,						hidden_size=32   ,						intermediate_size=37   ,						layer_norm_eps=1E-0_5   ,						num_attention_heads=4   ,						num_hidden_layers=5   ,						pad_token_id=1   ,						vocab_size=1000   ,						)
           return CLIPTextModel(SCREAMING_SNAKE_CASE_	)

       @property
       def        snake_case__     (       self	)       ->   Dict:

           """simple docstring"""



           def extract(*lowerCamelCase__   ,						**lowerCamelCase__	):






               class 			SCREAMING_SNAKE_CASE:

                      def __init__(       self	)       ->   int:

                          """simple docstring"""



                          __lowercase     							=		torch.ones([0]	)




                      def        snake_case__     (       self   ,						lowerCamelCase__	)       ->   str:

                          """simple docstring"""



                          self.pixel_values.to(SCREAMING_SNAKE_CASE_	)
                          return self

               return Out()

           return extract

       def        snake_case__     (       self	)       ->   int:

           """simple docstring"""



           __lowercase     							=		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
           __lowercase     							=		self.dummy_cond_unet
           __lowercase     							=		DDIMScheduler(
               beta_start=0.0_00_85   ,						beta_end=0.0_12   ,						beta_schedule="""scaled_linear"""   ,						clip_sample=SCREAMING_SNAKE_CASE_   ,						set_alpha_to_one=SCREAMING_SNAKE_CASE_   ,						)

           __lowercase     							=		self.dummy_vae
           __lowercase     							=		self.dummy_text_encoder
           __lowercase     							=		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)

           # make sure here that pndm scheduler skips prk
           __lowercase     							=		StableDiffusionPipeline(
               unet=SCREAMING_SNAKE_CASE_   ,						scheduler=SCREAMING_SNAKE_CASE_   ,						vae=SCREAMING_SNAKE_CASE_   ,						text_encoder=SCREAMING_SNAKE_CASE_   ,						tokenizer=SCREAMING_SNAKE_CASE_   ,						safety_checker=SCREAMING_SNAKE_CASE_   ,						feature_extractor=self.dummy_extractor   ,						)
           __lowercase     							=		sd_pipe.to(SCREAMING_SNAKE_CASE_	)
           sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

           __lowercase     							=		'''A painting of a squirrel eating a burger'''

           __lowercase     							=		torch.Generator(device=SCREAMING_SNAKE_CASE_	).manual_seed(0	)
           __lowercase     							=		sd_pipe([prompt]   ,						generator=SCREAMING_SNAKE_CASE_   ,						guidance_scale=6.0   ,						num_inference_steps=2   ,						output_type="""np"""	)
           __lowercase     							=		output.images

           __lowercase     							=		torch.Generator(device=SCREAMING_SNAKE_CASE_	).manual_seed(0	)
           __lowercase     							=		sd_pipe(
               [prompt]   ,						generator=SCREAMING_SNAKE_CASE_   ,						guidance_scale=6.0   ,						num_inference_steps=2   ,						output_type="""np"""   ,						return_dict=SCREAMING_SNAKE_CASE_   ,						)[0]

           __lowercase     							=		image[0, -3:, -3:, -1]
           __lowercase     							=		image_from_tuple[0, -3:, -3:, -1]

           assert image.shape == (1, 64, 64, 3)
           __lowercase     							=		np.array([0.57_53, 0.61_14, 0.50_01, 0.50_34, 0.54_70, 0.47_29, 0.49_71, 0.48_67, 0.48_67]	)

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

       def        snake_case__     (       self	)       ->   int:

           """simple docstring"""



           __lowercase     							=		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
           __lowercase     							=		self.dummy_cond_unet
           __lowercase     							=		PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE_	)
           __lowercase     							=		self.dummy_vae
           __lowercase     							=		self.dummy_text_encoder
           __lowercase     							=		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)

           # make sure here that pndm scheduler skips prk
           __lowercase     							=		StableDiffusionPipeline(
               unet=SCREAMING_SNAKE_CASE_   ,						scheduler=SCREAMING_SNAKE_CASE_   ,						vae=SCREAMING_SNAKE_CASE_   ,						text_encoder=SCREAMING_SNAKE_CASE_   ,						tokenizer=SCREAMING_SNAKE_CASE_   ,						safety_checker=SCREAMING_SNAKE_CASE_   ,						feature_extractor=self.dummy_extractor   ,						)
           __lowercase     							=		sd_pipe.to(SCREAMING_SNAKE_CASE_	)
           sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

           __lowercase     							=		'''A painting of a squirrel eating a burger'''
           __lowercase     							=		torch.Generator(device=SCREAMING_SNAKE_CASE_	).manual_seed(0	)
           __lowercase     							=		sd_pipe([prompt]   ,						generator=SCREAMING_SNAKE_CASE_   ,						guidance_scale=6.0   ,						num_inference_steps=2   ,						output_type="""np"""	)

           __lowercase     							=		output.images

           __lowercase     							=		torch.Generator(device=SCREAMING_SNAKE_CASE_	).manual_seed(0	)
           __lowercase     							=		sd_pipe(
               [prompt]   ,						generator=SCREAMING_SNAKE_CASE_   ,						guidance_scale=6.0   ,						num_inference_steps=2   ,						output_type="""np"""   ,						return_dict=SCREAMING_SNAKE_CASE_   ,						)[0]

           __lowercase     							=		image[0, -3:, -3:, -1]
           __lowercase     							=		image_from_tuple[0, -3:, -3:, -1]

           assert image.shape == (1, 64, 64, 3)
           __lowercase     							=		np.array([0.51_22, 0.57_12, 0.48_25, 0.50_53, 0.56_46, 0.47_69, 0.51_79, 0.48_94, 0.49_94]	)

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

       def        snake_case__     (       self	)       ->   Dict:

           """simple docstring"""



           __lowercase     							=		StableDiffusionPipeline.from_pretrained(
               """hf-internal-testing/tiny-stable-diffusion-lms-pipe"""   ,						safety_checker=SCREAMING_SNAKE_CASE_	)
           assert isinstance(SCREAMING_SNAKE_CASE_   ,						SCREAMING_SNAKE_CASE_	)
           assert isinstance(pipe.scheduler   ,						SCREAMING_SNAKE_CASE_	)
           assert pipe.safety_checker is None

           __lowercase     							=		pipe("""example prompt"""   ,						num_inference_steps=2	).images[0]
           assert image is not None

           # check that there's no error when saving a pipeline with one of the models being None
           with tempfile.TemporaryDirectory() as tmpdirname:
               pipe.save_pretrained(SCREAMING_SNAKE_CASE_	)
               __lowercase     							=		StableDiffusionPipeline.from_pretrained(SCREAMING_SNAKE_CASE_	)

           # sanity check that the pipeline still works
           assert pipe.safety_checker is None
           __lowercase     							=		pipe("""example prompt"""   ,						num_inference_steps=2	).images[0]
           assert image is not None




       @unittest.skipIf(torch_device != """cuda"""   ,						"""This test requires a GPU"""	)
       def        snake_case__     (       self	)       ->   Optional[Any]:

           """simple docstring"""



           __lowercase     							=		self.dummy_cond_unet
           __lowercase     							=		PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE_	)
           __lowercase     							=		self.dummy_vae
           __lowercase     							=		self.dummy_text_encoder
           __lowercase     							=		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""	)

           # put models in fp16
           __lowercase     							=		unet.half()
           __lowercase     							=		vae.half()
           __lowercase     							=		bert.half()

           # make sure here that pndm scheduler skips prk
           __lowercase     							=		StableDiffusionPipeline(
               unet=SCREAMING_SNAKE_CASE_   ,						scheduler=SCREAMING_SNAKE_CASE_   ,						vae=SCREAMING_SNAKE_CASE_   ,						text_encoder=SCREAMING_SNAKE_CASE_   ,						tokenizer=SCREAMING_SNAKE_CASE_   ,						safety_checker=SCREAMING_SNAKE_CASE_   ,						feature_extractor=self.dummy_extractor   ,						)
           __lowercase     							=		sd_pipe.to(SCREAMING_SNAKE_CASE_	)
           sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

           __lowercase     							=		'''A painting of a squirrel eating a burger'''
           __lowercase     							=		sd_pipe([prompt]   ,						num_inference_steps=2   ,						output_type="""np"""	).images

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






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

       def        snake_case__     (       self	)       ->   int:

           """simple docstring"""



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

       def        snake_case__     (       self	)       ->   List[Any]:

           """simple docstring"""



           __lowercase     							=		'''cuda'''
           __lowercase     							=		StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""	)
           __lowercase     							=		pipe.to(SCREAMING_SNAKE_CASE_	)
           pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

           __lowercase     							=		'''a photo of a cat'''
           __lowercase     							=		{
               '''editing_prompt''': ['''sunglasses'''],
               '''reverse_editing_direction''': [False],
               '''edit_warmup_steps''': 10,
               '''edit_guidance_scale''': 6,
               '''edit_threshold''': 0.95,
               '''edit_momentum_scale''': 0.5,
               '''edit_mom_beta''': 0.6,
           }

           __lowercase     							=		3
           __lowercase     							=		7

           # no sega enabled
           __lowercase     							=		torch.Generator(SCREAMING_SNAKE_CASE_	)
           generator.manual_seed(SCREAMING_SNAKE_CASE_	)
           __lowercase     							=		pipe(
               [prompt]   ,						generator=SCREAMING_SNAKE_CASE_   ,						guidance_scale=SCREAMING_SNAKE_CASE_   ,						num_inference_steps=50   ,						output_type="""np"""   ,						width=512   ,						height=512   ,						)

           __lowercase     							=		output.images
           __lowercase     							=		image[0, -3:, -3:, -1]
           __lowercase     							=		[
               0.34_67_31_13,
               0.38_49_27_33,
               0.37_59_73_52,
               0.34_08_63_35,
               0.35_65_07_48,
               0.35_57_92_05,
               0.3_38_47_63,
               0.34_34_02_36,
               0.3_57_32_71,
           ]

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

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

           # with sega enabled
           # generator = torch.manual_seed(seed)
           generator.manual_seed(SCREAMING_SNAKE_CASE_	)
           __lowercase     							=		pipe(
               [prompt]   ,						generator=SCREAMING_SNAKE_CASE_   ,						guidance_scale=SCREAMING_SNAKE_CASE_   ,						num_inference_steps=50   ,						output_type="""np"""   ,						width=512   ,						height=512   ,						**SCREAMING_SNAKE_CASE_   ,						)

           __lowercase     							=		output.images
           __lowercase     							=		image[0, -3:, -3:, -1]
           __lowercase     							=		[
               0.41_88_78_26,
               0.37_72_87_66,
               0.30_13_82_72,
               0.41_41_63_35,
               0.41_66_49_85,
               0.36_28_33_92,
               0.36_19_12_46,
               0.43_36_44_65,
               0.43_00_17_32,
           ]

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

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

       def        snake_case__     (       self	)       ->   Any:

           """simple docstring"""



           __lowercase     							=		'''cuda'''
           __lowercase     							=		StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""	)
           __lowercase     							=		pipe.to(SCREAMING_SNAKE_CASE_	)
           pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

           __lowercase     							=		'''an image of a crowded boulevard, realistic, 4k'''
           __lowercase     							=		{
               '''editing_prompt''': '''crowd, crowded, people''',
               '''reverse_editing_direction''': True,
               '''edit_warmup_steps''': 10,
               '''edit_guidance_scale''': 8.3,
               '''edit_threshold''': 0.9,
               '''edit_momentum_scale''': 0.5,
               '''edit_mom_beta''': 0.6,
           }

           __lowercase     							=		9
           __lowercase     							=		7

           # no sega enabled
           __lowercase     							=		torch.Generator(SCREAMING_SNAKE_CASE_	)
           generator.manual_seed(SCREAMING_SNAKE_CASE_	)
           __lowercase     							=		pipe(
               [prompt]   ,						generator=SCREAMING_SNAKE_CASE_   ,						guidance_scale=SCREAMING_SNAKE_CASE_   ,						num_inference_steps=50   ,						output_type="""np"""   ,						width=512   ,						height=512   ,						)

           __lowercase     							=		output.images
           __lowercase     							=		image[0, -3:, -3:, -1]
           __lowercase     							=		[
               0.43_49_79_98,
               0.91_81_40_65,
               0.7_54_07_39,
               0.55_58_02_05,
               0.8_46_72_65,
               0.5_38_96_91,
               0.62_57_45_06,
               0.58_89_77_63,
               0.50_92_67_57,
           ]

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

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

           # with sega enabled
           # generator = torch.manual_seed(seed)
           generator.manual_seed(SCREAMING_SNAKE_CASE_	)
           __lowercase     							=		pipe(
               [prompt]   ,						generator=SCREAMING_SNAKE_CASE_   ,						guidance_scale=SCREAMING_SNAKE_CASE_   ,						num_inference_steps=50   ,						output_type="""np"""   ,						width=512   ,						height=512   ,						**SCREAMING_SNAKE_CASE_   ,						)

           __lowercase     							=		output.images
           __lowercase     							=		image[0, -3:, -3:, -1]
           __lowercase     							=		[
               0.3_08_97_19,
               0.30_50_01_44,
               0.29_01_60_42,
               0.30_63_09_64,
               0.32_56_87,
               0.29_41_92_25,
               0.2_90_80_91,
               0.28_72_35_98,
               0.27_69_62_94,
           ]

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

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

       def        snake_case__     (       self	)       ->   Tuple:

           """simple docstring"""



           __lowercase     							=		'''cuda'''
           __lowercase     							=		StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""	)
           __lowercase     							=		pipe.to(SCREAMING_SNAKE_CASE_	)
           pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

           __lowercase     							=		'''a castle next to a river'''
           __lowercase     							=		{
               '''editing_prompt''': ['''boat on a river, boat''', '''monet, impression, sunrise'''],
               '''reverse_editing_direction''': False,
               '''edit_warmup_steps''': [15, 18],
               '''edit_guidance_scale''': 6,
               '''edit_threshold''': [0.9, 0.8],
               '''edit_momentum_scale''': 0.5,
               '''edit_mom_beta''': 0.6,
           }

           __lowercase     							=		48
           __lowercase     							=		7

           # no sega enabled
           __lowercase     							=		torch.Generator(SCREAMING_SNAKE_CASE_	)
           generator.manual_seed(SCREAMING_SNAKE_CASE_	)
           __lowercase     							=		pipe(
               [prompt]   ,						generator=SCREAMING_SNAKE_CASE_   ,						guidance_scale=SCREAMING_SNAKE_CASE_   ,						num_inference_steps=50   ,						output_type="""np"""   ,						width=512   ,						height=512   ,						)

           __lowercase     							=		output.images
           __lowercase     							=		image[0, -3:, -3:, -1]
           __lowercase     							=		[
               0.75_16_35_55,
               0.76_03_71_45,
               0.6_17_85,
               0.9_18_96_73,
               0.8_62_77_01,
               0.85_18_96_94,
               0.8_51_28_13,
               0.87_01_20_76,
               0.8_31_28_57,
           ]

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

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

           # with sega enabled
           # generator = torch.manual_seed(seed)
           generator.manual_seed(SCREAMING_SNAKE_CASE_	)
           __lowercase     							=		pipe(
               [prompt]   ,						generator=SCREAMING_SNAKE_CASE_   ,						guidance_scale=SCREAMING_SNAKE_CASE_   ,						num_inference_steps=50   ,						output_type="""np"""   ,						width=512   ,						height=512   ,						**SCREAMING_SNAKE_CASE_   ,						)

           __lowercase     							=		output.images
           __lowercase     							=		image[0, -3:, -3:, -1]
           __lowercase     							=		[
               0.73_55_33_65,
               0.7_53_72_71,
               0.74_34_19_05,
               0.66_48_03_56,
               0.6_47_29_25,
               0.63_03_94_16,
               0.64_81_29_05,
               0.6_74_97_17,
               0.6_51_71_02,
           ]

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

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




       def        snake_case__     (       self	)       ->   List[Any]:

           """simple docstring"""



           __lowercase     							=		'''cuda'''
           __lowercase     							=		StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""   ,						torch_dtype=torch.floataa	)
           __lowercase     							=		pipe.to(SCREAMING_SNAKE_CASE_	)
           pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

           __lowercase     							=		'''a photo of a cat'''
           __lowercase     							=		{
               '''editing_prompt''': ['''sunglasses'''],
               '''reverse_editing_direction''': [False],
               '''edit_warmup_steps''': 10,
               '''edit_guidance_scale''': 6,
               '''edit_threshold''': 0.95,
               '''edit_momentum_scale''': 0.5,
               '''edit_mom_beta''': 0.6,
           }

           __lowercase     							=		3
           __lowercase     							=		7

           # no sega enabled
           __lowercase     							=		torch.Generator(SCREAMING_SNAKE_CASE_	)
           generator.manual_seed(SCREAMING_SNAKE_CASE_	)
           __lowercase     							=		pipe(
               [prompt]   ,						generator=SCREAMING_SNAKE_CASE_   ,						guidance_scale=SCREAMING_SNAKE_CASE_   ,						num_inference_steps=50   ,						output_type="""np"""   ,						width=512   ,						height=512   ,						)

           __lowercase     							=		output.images
           __lowercase     							=		image[0, -3:, -3:, -1]
           __lowercase     							=		[
               0.34_88_76_95,
               0.3_87_69_53,
               0.3_75,
               0.34_42_38_28,
               0.3_58_15_43,
               0.35_71_77_73,
               0.3_38_37_89,
               0.34_57_03_12,
               0.35_93_75,
           ]

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

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

           # with sega enabled
           # generator = torch.manual_seed(seed)
           generator.manual_seed(SCREAMING_SNAKE_CASE_	)
           __lowercase     							=		pipe(
               [prompt]   ,						generator=SCREAMING_SNAKE_CASE_   ,						guidance_scale=SCREAMING_SNAKE_CASE_   ,						num_inference_steps=50   ,						output_type="""np"""   ,						width=512   ,						height=512   ,						**SCREAMING_SNAKE_CASE_   ,						)

           __lowercase     							=		output.images
           __lowercase     							=		image[0, -3:, -3:, -1]
           __lowercase     							=		[
               0.42_28_51_56,
               0.36_91_40_62,
               0.29_07_71_48,
               0.42_04_10_16,
               0.41_91_89_45,
               0.35_49_80_47,
               0.3_61_81_64,
               0.4_42_38_28,
               0.43_11_52_34,
           ]

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

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










'''simple docstring'''


import unittest

from transformers import YosoConfig, 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
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         YosoForMaskedLM,
         YosoForMultipleChoice,
         YosoForQuestionAnswering,
         YosoForSequenceClassification,
         YosoForTokenClassification,
         YosoModel,
     )
     from transformers.models.yoso.modeling_yoso import YOSO_PRETRAINED_MODEL_ARCHIVE_LIST






class 						lowerCAmelCase_	:







  '''simple docstring'''

  def __init__( self				:							List[Any]					,  SCREAMING_SNAKE_CASE_				:							Dict					,  SCREAMING_SNAKE_CASE_				:							int=13					,  SCREAMING_SNAKE_CASE_				:							List[str]=7					,  SCREAMING_SNAKE_CASE_				:							Union[str, Any]=True					,  SCREAMING_SNAKE_CASE_				:							Tuple=True					,  SCREAMING_SNAKE_CASE_				:							Optional[int]=True					,  SCREAMING_SNAKE_CASE_				:							Optional[int]=True					,  SCREAMING_SNAKE_CASE_				:							Dict=99					,  SCREAMING_SNAKE_CASE_				:							Tuple=32					,  SCREAMING_SNAKE_CASE_				:							List[Any]=5					,  SCREAMING_SNAKE_CASE_				:							List[Any]=4					,  SCREAMING_SNAKE_CASE_				:							Dict=37					,  SCREAMING_SNAKE_CASE_				:							str="gelu"					,  SCREAMING_SNAKE_CASE_				:							str=0.1					,  SCREAMING_SNAKE_CASE_				:							Optional[Any]=0.1					,  SCREAMING_SNAKE_CASE_				:							Dict=5_12					,  SCREAMING_SNAKE_CASE_				:							Union[str, Any]=16					,  SCREAMING_SNAKE_CASE_				:							Union[str, Any]=2					,  SCREAMING_SNAKE_CASE_				:							Any=0.02					,  SCREAMING_SNAKE_CASE_				:							Dict=3					,  SCREAMING_SNAKE_CASE_				:							Any=4					,  SCREAMING_SNAKE_CASE_				:							List[str]=None					,  )			->		Optional[int]:







       '''simple docstring'''



       A:		int          =		parent
       A:		Any          =		batch_size
       A:		List[Any]          =		seq_length
       A:		Union[str, Any]          =		is_training
       A:		List[str]          =		use_input_mask
       A:		List[str]          =		use_token_type_ids
       A:		Union[str, Any]          =		use_labels
       A:		str          =		vocab_size
       A:		Union[str, Any]          =		hidden_size
       A:		Tuple          =		num_hidden_layers
       A:		Dict          =		num_attention_heads
       A:		Any          =		intermediate_size
       A:		Tuple          =		hidden_act
       A:		Dict          =		hidden_dropout_prob
       A:		List[Any]          =		attention_probs_dropout_prob
       A:		int          =		max_position_embeddings
       A:		Any          =		type_vocab_size
       A:		int          =		type_sequence_label_size
       A:		str          =		initializer_range
       A:		Dict          =		num_labels
       A:		int          =		num_choices
       A:		str          =		scope



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







       '''simple docstring'''



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

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

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

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

       A:		List[str]          =		self.get_config()

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



  def 					_snake_case     ( self				:							Union[str, Any]	)			->		Dict:







       '''simple docstring'''



       return YosoConfig(
           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=SCREAMING_SNAKE_CASE_					,  initializer_range=self.initializer_range					,  )



  def 					_snake_case     ( self				:							Optional[Any]	)			->		Tuple:







       '''simple docstring'''



       A:		str          =		self.get_config()
       A:		Tuple          =		3_00
       return config



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







       '''simple docstring'''



       (
           (
           A
       )     ,    (
           A
       )     ,    (
           A
       )     ,    (
           A
       )     ,    (
           A
       )     ,    (
           A
       )     ,    (
           A
       )     ,    
       ):		Optional[int]          =		self.prepare_config_and_inputs()

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







       '''simple docstring'''



       A:		int          =		YosoModel(config=SCREAMING_SNAKE_CASE_	)
       model.to(SCREAMING_SNAKE_CASE_	)
       model.eval()
       A:		str          =		model(SCREAMING_SNAKE_CASE_					,  attention_mask=SCREAMING_SNAKE_CASE_					,  token_type_ids=SCREAMING_SNAKE_CASE_	)
       A:		str          =		model(SCREAMING_SNAKE_CASE_					,  token_type_ids=SCREAMING_SNAKE_CASE_	)
       A:		Any          =		model(SCREAMING_SNAKE_CASE_	)
       self.parent.assertEqual(result.last_hidden_state.shape					,  (self.batch_size, self.seq_length, self.hidden_size)	)



  def 					_snake_case     ( self				:							Optional[int]					,  SCREAMING_SNAKE_CASE_				:							List[str]					,  SCREAMING_SNAKE_CASE_				:							Any					,  SCREAMING_SNAKE_CASE_				:							Dict					,  SCREAMING_SNAKE_CASE_				:							List[str]					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							Optional[int]					,  SCREAMING_SNAKE_CASE_				:							int					,  SCREAMING_SNAKE_CASE_				:							Optional[Any]					,  SCREAMING_SNAKE_CASE_				:							List[Any]					,  )			->		List[Any]:







       '''simple docstring'''



       A:		int          =		True
       A:		int          =		YosoModel(SCREAMING_SNAKE_CASE_	)
       model.to(SCREAMING_SNAKE_CASE_	)
       model.eval()
       A:		Optional[int]          =		model(
           SCREAMING_SNAKE_CASE_					,  attention_mask=SCREAMING_SNAKE_CASE_					,  token_type_ids=SCREAMING_SNAKE_CASE_					,  encoder_hidden_states=SCREAMING_SNAKE_CASE_					,  encoder_attention_mask=SCREAMING_SNAKE_CASE_					,  )
       A:		Any          =		model(
           SCREAMING_SNAKE_CASE_					,  attention_mask=SCREAMING_SNAKE_CASE_					,  token_type_ids=SCREAMING_SNAKE_CASE_					,  encoder_hidden_states=SCREAMING_SNAKE_CASE_					,  )
       A:		List[str]          =		model(SCREAMING_SNAKE_CASE_					,  attention_mask=SCREAMING_SNAKE_CASE_					,  token_type_ids=SCREAMING_SNAKE_CASE_	)
       self.parent.assertEqual(result.last_hidden_state.shape					,  (self.batch_size, self.seq_length, self.hidden_size)	)



  def 					_snake_case     ( self				:							Optional[int]					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							Tuple					,  SCREAMING_SNAKE_CASE_				:							Dict					,  SCREAMING_SNAKE_CASE_				:							Dict					,  SCREAMING_SNAKE_CASE_				:							Optional[Any]					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							Tuple	)			->		Union[str, Any]:







       '''simple docstring'''



       A:		List[Any]          =		YosoForMaskedLM(config=SCREAMING_SNAKE_CASE_	)
       model.to(SCREAMING_SNAKE_CASE_	)
       model.eval()
       A:		Any          =		model(SCREAMING_SNAKE_CASE_					,  attention_mask=SCREAMING_SNAKE_CASE_					,  token_type_ids=SCREAMING_SNAKE_CASE_					,  labels=SCREAMING_SNAKE_CASE_	)
       self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.seq_length, self.vocab_size)	)



  def 					_snake_case     ( self				:							Dict					,  SCREAMING_SNAKE_CASE_				:							Optional[Any]					,  SCREAMING_SNAKE_CASE_				:							Any					,  SCREAMING_SNAKE_CASE_				:							Optional[int]					,  SCREAMING_SNAKE_CASE_				:							Tuple					,  SCREAMING_SNAKE_CASE_				:							Dict					,  SCREAMING_SNAKE_CASE_				:							List[Any]					,  SCREAMING_SNAKE_CASE_				:							Union[str, Any]	)			->		int:







       '''simple docstring'''



       A:		Optional[int]          =		YosoForQuestionAnswering(config=SCREAMING_SNAKE_CASE_	)
       model.to(SCREAMING_SNAKE_CASE_	)
       model.eval()
       A:		Dict          =		model(
           SCREAMING_SNAKE_CASE_					,  attention_mask=SCREAMING_SNAKE_CASE_					,  token_type_ids=SCREAMING_SNAKE_CASE_					,  start_positions=SCREAMING_SNAKE_CASE_					,  end_positions=SCREAMING_SNAKE_CASE_					,  )
       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 					_snake_case     ( self				:							Optional[int]					,  SCREAMING_SNAKE_CASE_				:							List[Any]					,  SCREAMING_SNAKE_CASE_				:							List[str]					,  SCREAMING_SNAKE_CASE_				:							Tuple					,  SCREAMING_SNAKE_CASE_				:							Dict					,  SCREAMING_SNAKE_CASE_				:							Dict					,  SCREAMING_SNAKE_CASE_				:							Tuple					,  SCREAMING_SNAKE_CASE_				:							Dict	)			->		List[Any]:







       '''simple docstring'''



       A:		str          =		self.num_labels
       A:		int          =		YosoForSequenceClassification(SCREAMING_SNAKE_CASE_	)
       model.to(SCREAMING_SNAKE_CASE_	)
       model.eval()
       A:		Any          =		model(SCREAMING_SNAKE_CASE_					,  attention_mask=SCREAMING_SNAKE_CASE_					,  token_type_ids=SCREAMING_SNAKE_CASE_					,  labels=SCREAMING_SNAKE_CASE_	)
       self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.num_labels)	)



  def 					_snake_case     ( self				:							Dict					,  SCREAMING_SNAKE_CASE_				:							Any					,  SCREAMING_SNAKE_CASE_				:							Dict					,  SCREAMING_SNAKE_CASE_				:							List[str]					,  SCREAMING_SNAKE_CASE_				:							Any					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							Dict					,  SCREAMING_SNAKE_CASE_				:							Dict	)			->		Optional[Any]:







       '''simple docstring'''



       A:		Tuple          =		self.num_labels
       A:		Optional[Any]          =		YosoForTokenClassification(config=SCREAMING_SNAKE_CASE_	)
       model.to(SCREAMING_SNAKE_CASE_	)
       model.eval()
       A:		Dict          =		model(SCREAMING_SNAKE_CASE_					,  attention_mask=SCREAMING_SNAKE_CASE_					,  token_type_ids=SCREAMING_SNAKE_CASE_					,  labels=SCREAMING_SNAKE_CASE_	)
       self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.seq_length, self.num_labels)	)



  def 					_snake_case     ( self				:							Union[str, Any]					,  SCREAMING_SNAKE_CASE_				:							List[str]					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							Optional[int]					,  SCREAMING_SNAKE_CASE_				:							Optional[Any]					,  SCREAMING_SNAKE_CASE_				:							Any					,  SCREAMING_SNAKE_CASE_				:							str					,  SCREAMING_SNAKE_CASE_				:							Optional[Any]	)			->		Any:







       '''simple docstring'''



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




  def 					_snake_case     ( self				:							int	)			->		Optional[int]:







       '''simple docstring'''



       A:		List[Any]          =		self.prepare_config_and_inputs()
       (
           (
           A
       )     ,    (
           A
       )     ,    (
           A
       )     ,    (
           A
       )     ,    (
           A
       )     ,    (
           A
       )     ,    (
           A
       )     ,    
       ):		Optional[int]          =		config_and_inputs
       A:		Any          =		{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
       return config, inputs_dict






@require_torch
class 						lowerCAmelCase_	(       UpperCAmelCase_       ,					UpperCAmelCase_       ,					unittest.TestCase							):







  '''simple docstring'''

  UpperCamelCase_			:      List[str]										=				(
      (
          YosoModel,
          YosoForMaskedLM,
          YosoForMultipleChoice,
          YosoForQuestionAnswering,
          YosoForSequenceClassification,
          YosoForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  UpperCamelCase_			:      Tuple										=				False
  UpperCamelCase_			:      Tuple										=				False
  UpperCamelCase_			:      List[str]										=				False

  UpperCamelCase_			:      List[str]										=				()
  UpperCamelCase_			:      Optional[int]										=				(
      {
          """feature-extraction""": YosoModel,
          """fill-mask""": YosoForMaskedLM,
          """question-answering""": YosoForQuestionAnswering,
          """text-classification""": YosoForSequenceClassification,
          """token-classification""": YosoForTokenClassification,
          """zero-shot""": YosoForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )



  def 					_snake_case     ( self				:							Dict	)			->		Any:







       '''simple docstring'''



       A:		Tuple          =		YosoModelTester(self	)
       A:		str          =		ConfigTester(self					,  config_class=SCREAMING_SNAKE_CASE_					,  hidden_size=37	)



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







       '''simple docstring'''



       self.config_tester.run_common_tests()



  def 					_snake_case     ( self				:							str	)			->		Optional[Any]:







       '''simple docstring'''



       A:		int          =		self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_	)



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







       '''simple docstring'''



       A:		Any          =		self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
            A:		Tuple          =		type
            self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_	)



  def 					_snake_case     ( self				:							Tuple	)			->		Optional[int]:







       '''simple docstring'''



       A:		List[str]          =		self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_	)



  def 					_snake_case     ( self				:							List[Any]	)			->		str:







       '''simple docstring'''



       A:		str          =		self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE_	)



  def 					_snake_case     ( self				:							List[str]	)			->		str:







       '''simple docstring'''



       A:		Tuple          =		self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE_	)



  def 					_snake_case     ( self				:							Tuple	)			->		Dict:







       '''simple docstring'''



       A:		int          =		self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE_	)



  def 					_snake_case     ( self				:							Dict	)			->		Dict:







       '''simple docstring'''



       A:		Dict          =		self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_	)



  @slow
  def 					_snake_case     ( self				:							int	)			->		int:







       '''simple docstring'''



       for model_name in YOSO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            A:		Dict          =		YosoModel.from_pretrained(SCREAMING_SNAKE_CASE_	)
            self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)




  def 					_snake_case     ( self				:							Optional[Any]	)			->		Any:







       '''simple docstring'''



       return






@require_torch
class 						lowerCAmelCase_	(       unittest.TestCase							):







  '''simple docstring'''

  @slow
  def 					_snake_case     ( self				:							Union[str, Any]	)			->		Union[str, Any]:







       '''simple docstring'''



       A:		List[Any]          =		YosoModel.from_pretrained('''uw-madison/yoso-4096'''	)
       A:		Dict          =		torch.tensor([[0, 1, 2, 3, 4, 5]]	)

       with torch.no_grad():
            A:		Union[str, Any]          =		model(SCREAMING_SNAKE_CASE_	)[0]

       A:		List[str]          =		torch.Size((1, 6, 7_68)	)
       self.assertEqual(output.shape					,  SCREAMING_SNAKE_CASE_	)

       A:		int          =		torch.tensor(
           [[[-0.0611, 0.1242, 0.0840], [0.0280, -0.0048, 0.1125], [0.0106, 0.0226, 0.0751]]]	)

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



  @slow
  def 					_snake_case     ( self				:							int	)			->		Tuple:







       '''simple docstring'''



       A:		List[str]          =		YosoForMaskedLM.from_pretrained('''uw-madison/yoso-4096'''	)
       A:		Optional[int]          =		torch.tensor([[0, 1, 2, 3, 4, 5]]	)

       with torch.no_grad():
            A:		str          =		model(SCREAMING_SNAKE_CASE_	)[0]

       A:		Optional[int]          =		5_02_65

       A:		List[str]          =		torch.Size((1, 6, vocab_size)	)
       self.assertEqual(output.shape					,  SCREAMING_SNAKE_CASE_	)

       A:		Dict          =		torch.tensor(
           [[[-2.1313, -3.7285, -2.2407], [-2.7047, -3.3314, -2.6408], [0.0629, -2.5166, -0.3356]]]	)

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




  @slow
  def 					_snake_case     ( self				:							Dict	)			->		Tuple:







       '''simple docstring'''



       A:		List[str]          =		YosoForMaskedLM.from_pretrained('''uw-madison/yoso-4096'''	)
       A:		Optional[int]          =		torch.arange(40_96	).unsqueeze(0	)

       with torch.no_grad():
            A:		int          =		model(SCREAMING_SNAKE_CASE_	)[0]

       A:		Optional[int]          =		5_02_65

       A:		Dict          =		torch.Size((1, 40_96, vocab_size)	)
       self.assertEqual(output.shape					,  SCREAMING_SNAKE_CASE_	)

       A:		Optional[int]          =		torch.tensor(
           [[[-2.3914, -4.3742, -5.0956], [-4.0988, -4.2384, -7.0406], [-3.1427, -3.7192, -6.6800]]]	)

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






