






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

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, 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_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class    __SCREAMING_SNAKE_CASE					(	__SCREAMING_SNAKE_CASE		, __SCREAMING_SNAKE_CASE		, unittest.TestCase ):
 snake_case_     							=							StableDiffusionXLImgaImgPipeline
 snake_case_     							=							TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
 snake_case_     							=							PipelineTesterMixin.required_optional_params - {"""latents"""}
 snake_case_     							=							TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 snake_case_     							=							IMAGE_TO_IMAGE_IMAGE_PARAMS
 snake_case_     							=							IMAGE_TO_IMAGE_IMAGE_PARAMS




 def    __magic_name__		(					self		:							Dict					)		->  Optional[int]:
       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE__   : List[Any]													=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''')    ,					attention_head_dim=(2, 4)    ,					use_linear_projection=_SCREAMING_SNAKE_CASE    ,					addition_embed_type='''text_time'''    ,					addition_time_embed_dim=8    ,					transformer_layers_per_block=(1, 2)    ,					projection_class_embeddings_input_dim=80    ,					cross_attention_dim=64    ,					)
       SCREAMING_SNAKE_CASE__   : Tuple													=EulerDiscreteScheduler(
           beta_start=0.00085    ,					beta_end=0.012    ,					steps_offset=1    ,					beta_schedule='''scaled_linear'''    ,					timestep_spacing='''leading'''    ,					)
       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE__   : str													=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    ,					sample_size=1_28    ,					)
       torch.manual_seed(0					)
       SCREAMING_SNAKE_CASE__   : Optional[int]													=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    ,					hidden_act='''gelu'''    ,					projection_dim=32    ,					)
       SCREAMING_SNAKE_CASE__   : List[Any]													=CLIPTextModel(_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : List[Any]													=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    ,					local_files_only=_SCREAMING_SNAKE_CASE					)

       SCREAMING_SNAKE_CASE__   : Optional[Any]													=CLIPTextModelWithProjection(_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : int													=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    ,					local_files_only=_SCREAMING_SNAKE_CASE					)

       SCREAMING_SNAKE_CASE__   : Any													={
           "unet": unet,
           "scheduler": scheduler,
           "vae": vae,
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
           "text_encoder_2": text_encoder_a,
           "tokenizer_2": tokenizer_a,
           # "safety_checker": None,
           # "feature_extractor": None,
       }
       return components




 def    __magic_name__		(					self		:							int    ,					__lowercase		:							Dict    ,					__lowercase		:							Any=0					)		->  int:
       SCREAMING_SNAKE_CASE__   : int													=floats_tensor((1, 3, 32, 32)    ,					rng=random.Random(_SCREAMING_SNAKE_CASE					)					).to(_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : List[Any]													=image / 2 + 0.5
       if str(_SCREAMING_SNAKE_CASE					).startswith('''mps'''					):
             SCREAMING_SNAKE_CASE__   : int													=torch.manual_seed(_SCREAMING_SNAKE_CASE					)
       else:
             SCREAMING_SNAKE_CASE__   : Any													=torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : int													={
           "prompt": "A painting of a squirrel eating a burger",
           "image": image,
           "generator": generator,
           "num_inference_steps": 2,
           "guidance_scale": 5.0,
           "output_type": "numpy",
           "strength": 0.75,
       }
       return inputs




 def    __magic_name__		(					self		:							Tuple					)		->  Tuple:
       SCREAMING_SNAKE_CASE__   : List[Any]													="cpu"  # ensure determinism for the device-dependent torch.Generator
       SCREAMING_SNAKE_CASE__   : Optional[Any]													=self.get_dummy_components()
       SCREAMING_SNAKE_CASE__   : List[Any]													=StableDiffusionXLImgaImgPipeline(**_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : int													=sd_pipe.to(_SCREAMING_SNAKE_CASE					)
       sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

       SCREAMING_SNAKE_CASE__   : List[Any]													=self.get_dummy_inputs(_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : Union[str, Any]													=sd_pipe(**_SCREAMING_SNAKE_CASE					).images
       SCREAMING_SNAKE_CASE__   : str													=image[0, -3:, -3:, -1]

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

       SCREAMING_SNAKE_CASE__   : int													=np.array([0.4656, 0.4840, 0.4439, 0.6698, 0.5574, 0.4524, 0.5799, 0.5943, 0.5165]					)

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




 def    __magic_name__		(					self		:							Dict					)		->  Dict:
       super().test_attention_slicing_forward_pass(expected_max_diff=3e-3					)




 def    __magic_name__		(					self		:							Optional[int]					)		->  Union[str, Any]:
       super().test_inference_batch_single_identical(expected_max_diff=3e-3					)




 def    __magic_name__		(					self		:							List[str]					)		->  Optional[int]:
       pass




 def    __magic_name__		(					self		:							str					)		->  List[str]:
       SCREAMING_SNAKE_CASE__   : Dict													=self.get_dummy_components()
       SCREAMING_SNAKE_CASE__   : str													=StableDiffusionXLImgaImgPipeline(**_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : Dict													=sd_pipe.to(_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : int													=sd_pipe.to(_SCREAMING_SNAKE_CASE					)
       sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

       # forward without prompt embeds
       SCREAMING_SNAKE_CASE__   : str													=self.get_dummy_inputs(_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : str													=3 * ["this is a negative prompt"]
       SCREAMING_SNAKE_CASE__   : Union[str, Any]													=negative_prompt
       SCREAMING_SNAKE_CASE__   : Optional[int]													=3 * [inputs["prompt"]]

       SCREAMING_SNAKE_CASE__   : Dict													=sd_pipe(**_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : Union[str, Any]													=output.images[0, -3:, -3:, -1]

       # forward with prompt embeds
       SCREAMING_SNAKE_CASE__   : Union[str, Any]													=self.get_dummy_inputs(_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : Optional[Any]													=3 * ["this is a negative prompt"]
       SCREAMING_SNAKE_CASE__   : List[str]													=3 * [inputs.pop('''prompt'''					)]

       (
           SCREAMING_SNAKE_CASE__	
       )   : Optional[Any]													=sd_pipe.encode_prompt(_SCREAMING_SNAKE_CASE    ,					negative_prompt=_SCREAMING_SNAKE_CASE					)

       SCREAMING_SNAKE_CASE__   : str													=sd_pipe(
           **_SCREAMING_SNAKE_CASE    ,					prompt_embeds=_SCREAMING_SNAKE_CASE    ,					negative_prompt_embeds=_SCREAMING_SNAKE_CASE    ,					pooled_prompt_embeds=_SCREAMING_SNAKE_CASE    ,					negative_pooled_prompt_embeds=_SCREAMING_SNAKE_CASE    ,					)
       SCREAMING_SNAKE_CASE__   : Dict													=output.images[0, -3:, -3:, -1]

       # make sure that it's equal
       assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()					).max() < 1e-4






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




 def    __magic_name__		(					self		:							Optional[Any]					)		->  Any:
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




 def    __magic_name__		(					self		:							Optional[int]    ,					__lowercase		:							Dict    ,					__lowercase		:							str="cpu"    ,					__lowercase		:							Optional[int]=torch.floataa    ,					__lowercase		:							Optional[Any]=0					)		->  Optional[int]:
       SCREAMING_SNAKE_CASE__   : Tuple													=torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : Any													=np.random.RandomState(_SCREAMING_SNAKE_CASE					).standard_normal((1, 4, 64, 64)					)
       SCREAMING_SNAKE_CASE__   : Tuple													=torch.from_numpy(_SCREAMING_SNAKE_CASE					).to(device=_SCREAMING_SNAKE_CASE    ,					dtype=_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : Optional[int]													={
           "prompt": "a photograph of an astronaut riding a horse",
           "latents": latents,
           "generator": generator,
           "num_inference_steps": 3,
           "guidance_scale": 7.5,
           "output_type": "numpy",
       }
       return inputs




 def    __magic_name__		(					self		:							Union[str, Any]					)		->  Any:
       SCREAMING_SNAKE_CASE__   : Union[str, Any]													=DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''					)
       pipe.to(_SCREAMING_SNAKE_CASE					)
       pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

       SCREAMING_SNAKE_CASE__   : Dict													=self.get_inputs(_SCREAMING_SNAKE_CASE					)
       SCREAMING_SNAKE_CASE__   : Optional[Any]													=pipe(**_SCREAMING_SNAKE_CASE					).images
       SCREAMING_SNAKE_CASE__   : Optional[Any]													=image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 5_12, 5_12, 3)
       SCREAMING_SNAKE_CASE__   : Union[str, Any]													=np.array([0.49493, 0.47896, 0.40798, 0.54214, 0.53212, 0.48202, 0.47656, 0.46329, 0.48506]					)
       assert np.abs(image_slice - expected_slice					).max() < 7e-3




"""simple docstring"""






import unittest

from transformers import SqueezeBertConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, 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 (
	    SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
	    SqueezeBertForMaskedLM,
	    SqueezeBertForMultipleChoice,
	    SqueezeBertForQuestionAnswering,
	    SqueezeBertForSequenceClassification,
	    SqueezeBertForTokenClassification,
	    SqueezeBertModel,
	)



class 			A__   (  __SCREAMING_SNAKE_CASE							):




	'''simple docstring'''





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

				"""simple docstring"""

				__lowerCAmelCase      :  List[str]												=       parent
				__lowerCAmelCase      :  Optional[Any]												=       batch_size
				__lowerCAmelCase      :  Union[str, Any]												=       seq_length
				__lowerCAmelCase      :  Optional[Any]												=       is_training
				__lowerCAmelCase      :  Optional[int]												=       use_input_mask
				__lowerCAmelCase      :  Dict												=       use_token_type_ids
				__lowerCAmelCase      :  Dict												=       use_labels
				__lowerCAmelCase      :  Dict												=       vocab_size
				__lowerCAmelCase      :  Tuple												=       hidden_size
				__lowerCAmelCase      :  List[Any]												=       num_hidden_layers
				__lowerCAmelCase      :  Union[str, Any]												=       num_attention_heads
				__lowerCAmelCase      :  Tuple												=       intermediate_size
				__lowerCAmelCase      :  List[Any]												=       hidden_act
				__lowerCAmelCase      :  Optional[Any]												=       hidden_dropout_prob
				__lowerCAmelCase      :  Optional[Any]												=       attention_probs_dropout_prob
				__lowerCAmelCase      :  Optional[int]												=       max_position_embeddings
				__lowerCAmelCase      :  Union[str, Any]												=       type_vocab_size
				__lowerCAmelCase      :  Optional[int]												=       type_sequence_label_size
				__lowerCAmelCase      :  Dict												=       initializer_range
				__lowerCAmelCase      :  Tuple												=       num_labels
				__lowerCAmelCase      :  Optional[Any]												=       num_choices
				__lowerCAmelCase      :  Union[str, Any]												=       scope
				__lowerCAmelCase      :  Optional[Any]												=       q_groups
				__lowerCAmelCase      :  Optional[int]												=       k_groups
				__lowerCAmelCase      :  Any												=       v_groups
				__lowerCAmelCase      :  int												=       post_attention_groups
				__lowerCAmelCase      :  List[str]												=       intermediate_groups
				__lowerCAmelCase      :  Optional[Any]												=       output_groups

	def 							_SCREAMING_SNAKE_CASE   (     self:		Dict)				->       List[str]:

				"""simple docstring"""

				__lowerCAmelCase      :  str												=       ids_tensor([self.batch_size, self.seq_length]       ,			self.vocab_size)

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

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

				__lowerCAmelCase      :  Any												=       self.get_config()

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

	def 							_SCREAMING_SNAKE_CASE   (     self:		Tuple)				->       int:

				"""simple docstring"""

				return SqueezeBertConfig(
				    embedding_size=self.hidden_size       ,			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       ,			attention_probs_dropout_prob=self.hidden_dropout_prob       ,			attention_dropout=self.attention_probs_dropout_prob       ,			max_position_embeddings=self.max_position_embeddings       ,			initializer_range=self.initializer_range       ,			q_groups=self.q_groups       ,			k_groups=self.k_groups       ,			v_groups=self.v_groups       ,			post_attention_groups=self.post_attention_groups       ,			intermediate_groups=self.intermediate_groups       ,			output_groups=self.output_groups       ,			)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Tuple       ,			_SCREAMING_SNAKE_CASE:		List[str]       ,			_SCREAMING_SNAKE_CASE:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		Tuple       ,			_SCREAMING_SNAKE_CASE:		Dict       ,			_SCREAMING_SNAKE_CASE:		Dict       ,			_SCREAMING_SNAKE_CASE:		Tuple)				->       Tuple:

				"""simple docstring"""

				__lowerCAmelCase      :  List[Any]												=       SqueezeBertModel(config=_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()
				__lowerCAmelCase      :  List[Any]												=       model(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Dict												=       model(_SCREAMING_SNAKE_CASE)
				self.parent.assertEqual(result.last_hidden_state.shape       ,			(self.batch_size, self.seq_length, self.hidden_size))

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

				"""simple docstring"""

				__lowerCAmelCase      :  int												=       SqueezeBertForMaskedLM(config=_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()
				__lowerCAmelCase      :  Dict												=       model(_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE       ,			labels=_SCREAMING_SNAKE_CASE)
				self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.seq_length, self.vocab_size))

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

				"""simple docstring"""

				__lowerCAmelCase      :  str												=       SqueezeBertForQuestionAnswering(config=_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()
				__lowerCAmelCase      :  Union[str, Any]												=       model(
				    _SCREAMING_SNAKE_CASE       ,			attention_mask=_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 							_SCREAMING_SNAKE_CASE   (     self:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		Tuple       ,			_SCREAMING_SNAKE_CASE:		Any       ,			_SCREAMING_SNAKE_CASE:		Tuple)				->       Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  List[Any]												=       self.num_labels
				__lowerCAmelCase      :  Union[str, Any]												=       SqueezeBertForSequenceClassification(_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()
				__lowerCAmelCase      :  int												=       model(_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE       ,			labels=_SCREAMING_SNAKE_CASE)
				self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.num_labels))

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

				"""simple docstring"""

				__lowerCAmelCase      :  Dict												=       self.num_labels
				__lowerCAmelCase      :  Optional[int]												=       SqueezeBertForTokenClassification(config=_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()

				__lowerCAmelCase      :  List[str]												=       model(_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE       ,			labels=_SCREAMING_SNAKE_CASE)
				self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.seq_length, self.num_labels))

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

				"""simple docstring"""

				__lowerCAmelCase      :  List[str]												=       self.num_choices
				__lowerCAmelCase      :  str												=       SqueezeBertForMultipleChoice(config=_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()
				__lowerCAmelCase      :  int												=       input_ids.unsqueeze(1).expand(-1       ,			self.num_choices       ,			-1).contiguous()
				__lowerCAmelCase      :  Union[str, Any]												=       input_mask.unsqueeze(1).expand(-1       ,			self.num_choices       ,			-1).contiguous()
				__lowerCAmelCase      :  str												=       model(
				    _SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE       ,			labels=_SCREAMING_SNAKE_CASE       ,			)
				self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.num_choices))







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

				"""simple docstring"""

				__lowerCAmelCase      :  Optional[int]												=       self.prepare_config_and_inputs()
				((__lowerCAmelCase)						,	(__lowerCAmelCase)						,	(__lowerCAmelCase)						,	(__lowerCAmelCase)						,	(__lowerCAmelCase)						,	(__lowerCAmelCase))      :  Union[str, Any]												=       config_and_inputs
				__lowerCAmelCase      :  int												=       {"input_ids": input_ids, "attention_mask": input_mask}
				return config, inputs_dict






@require_torch
class 			A__   (  __SCREAMING_SNAKE_CASE  ,		__SCREAMING_SNAKE_CASE  ,		unittest.TestCase							):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE          =	(
	    (
	        SqueezeBertModel,
	        SqueezeBertForMaskedLM,
	        SqueezeBertForMultipleChoice,
	        SqueezeBertForQuestionAnswering,
	        SqueezeBertForSequenceClassification,
	        SqueezeBertForTokenClassification,
	    )
	    if is_torch_available()
	    else None
	)
	SCREAMING_SNAKE_CASE          =	(
	    {
	        'feature-extraction': SqueezeBertModel,
	        'fill-mask': SqueezeBertForMaskedLM,
	        'question-answering': SqueezeBertForQuestionAnswering,
	        'text-classification': SqueezeBertForSequenceClassification,
	        'token-classification': SqueezeBertForTokenClassification,
	        'zero-shot': SqueezeBertForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	SCREAMING_SNAKE_CASE          =	False
	SCREAMING_SNAKE_CASE          =	True
	SCREAMING_SNAKE_CASE          =	False

	def 							_SCREAMING_SNAKE_CASE   (     self:		Dict)				->       Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  Any												=       SqueezeBertModelTester(self)
				__lowerCAmelCase      :  Optional[int]												=       ConfigTester(self       ,			config_class=_SCREAMING_SNAKE_CASE       ,			dim=37)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Any)				->       Tuple:

				"""simple docstring"""

				self.config_tester.run_common_tests()

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

				"""simple docstring"""

				__lowerCAmelCase      :  Any												=       self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_squeezebert_model(*_SCREAMING_SNAKE_CASE)

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

				"""simple docstring"""

				__lowerCAmelCase      :  int												=       self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_squeezebert_for_masked_lm(*_SCREAMING_SNAKE_CASE)

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

				"""simple docstring"""

				__lowerCAmelCase      :  List[Any]												=       self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_squeezebert_for_question_answering(*_SCREAMING_SNAKE_CASE)

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

				"""simple docstring"""

				__lowerCAmelCase      :  Tuple												=       self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_squeezebert_for_sequence_classification(*_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Any)				->       int:

				"""simple docstring"""

				__lowerCAmelCase      :  List[str]												=       self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_squeezebert_for_token_classification(*_SCREAMING_SNAKE_CASE)

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

				"""simple docstring"""

				__lowerCAmelCase      :  Any												=       self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_squeezebert_for_multiple_choice(*_SCREAMING_SNAKE_CASE)







	@slow
	def 							_SCREAMING_SNAKE_CASE   (     self:		Any)				->       Dict:

				"""simple docstring"""

				for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							__lowerCAmelCase      :  Optional[Any]												=       SqueezeBertModel.from_pretrained(_SCREAMING_SNAKE_CASE)
							self.assertIsNotNone(_SCREAMING_SNAKE_CASE)






@require_sentencepiece
@require_tokenizers
@require_torch
class 			A__   (  unittest.TestCase							):




	'''simple docstring'''





	@slow
	def 							_SCREAMING_SNAKE_CASE   (     self:		int)				->       List[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  Optional[Any]												=       SqueezeBertForSequenceClassification.from_pretrained("squeezebert/squeezebert-mnli")

				__lowerCAmelCase      :  List[Any]												=       torch.tensor([[1, 2_9414, 232, 328, 740, 1140, 1_2695, 69, 13, 1588, 2]])
				__lowerCAmelCase      :  List[Any]												=       model(_SCREAMING_SNAKE_CASE)[0]
				__lowerCAmelCase      :  Any												=       torch.Size((1, 3))
				self.assertEqual(output.shape       ,			_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Union[str, Any]												=       torch.tensor([[0.6401, -0.0349, -0.6041]])
				self.assertTrue(torch.allclose(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE       ,			atol=1e-4))