import gc
import random
import tempfile
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMInverseScheduler,
    DDIMScheduler,
    DPMSolverMultistepInverseScheduler,
    DPMSolverMultistepScheduler,
    StableDiffusionDiffEditPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_image, slow
from diffusers.utils.testing_utils import enable_full_determinism, floats_tensor, require_torch_gpu, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class        lowerCamelCase__				(   UpperCAmelCase       ,UpperCAmelCase       ,unittest.TestCase				):
			UpperCamelCase__    							=StableDiffusionDiffEditPipeline
			UpperCamelCase__    							=TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"height", "width", "image"} | {"image_latents"}
			UpperCamelCase__    							=TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {"image"} | {"image_latents"}
			UpperCamelCase__    							=frozenset(
			    []				)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
			UpperCamelCase__    							=frozenset([]				)





			def   SCREAMING_SNAKE_CASE__     (       self	:	str						):
										torch.manual_seed(0						)
										_lowerCAmelCase							=							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  ,  attention_head_dim=(2, 4)  ,  use_linear_projection=lowercase__  ,  )
										_lowerCAmelCase							=							DDIMScheduler(
										    beta_start=0.0_0_0_8_5  ,  beta_end=0.0_1_2  ,  beta_schedule='scaled_linear'  ,  clip_sample=lowercase__  ,  set_alpha_to_one=lowercase__  ,  )
										_lowerCAmelCase							=							DDIMInverseScheduler(
										    beta_start=0.0_0_0_8_5  ,  beta_end=0.0_1_2  ,  beta_schedule='scaled_linear'  ,  clip_sample=lowercase__  ,  set_alpha_to_zero=lowercase__  ,  )
										torch.manual_seed(0						)
										_lowerCAmelCase							=							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						)
										_lowerCAmelCase							=							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=5_12  ,  )
										_lowerCAmelCase							=							CLIPTextModel(lowercase__						)
										_lowerCAmelCase							=							CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

										_lowerCAmelCase							=							{
										    'unet': unet,
										    'scheduler': scheduler,
										    'inverse_scheduler': inverse_scheduler,
										    'vae': vae,
										    'text_encoder': text_encoder,
										    'tokenizer': tokenizer,
										    'safety_checker': None,
										    'feature_extractor': None,
										}

										return components





			def   SCREAMING_SNAKE_CASE__     (       self	:	Tuple  ,  lowercase__	:	Any  ,  lowercase__	:	List[str]=0						):
										_lowerCAmelCase							=							floats_tensor((1, 16, 16)  ,  rng=random.Random(lowercase__						)						).to(lowercase__						)
										_lowerCAmelCase							=							floats_tensor((1, 2, 4, 16, 16)  ,  rng=random.Random(lowercase__						)						).to(lowercase__						)
										if str(lowercase__						).startswith('mps'						):
																	_lowerCAmelCase							=							torch.manual_seed(lowercase__						)
										else:
																	_lowerCAmelCase							=							torch.Generator(device=lowercase__						).manual_seed(lowercase__						)
										_lowerCAmelCase							=							{
										    'prompt': 'a dog and a newt',
										    'mask_image': mask,
										    'image_latents': latents,
										    'generator': generator,
										    'num_inference_steps': 2,
										    'inpaint_strength': 1.0,
										    'guidance_scale': 6.0,
										    'output_type': 'numpy',
										}

										return inputs





			def   SCREAMING_SNAKE_CASE__     (       self	:	str  ,  lowercase__	:	Tuple  ,  lowercase__	:	Dict=0						):
										_lowerCAmelCase							=							floats_tensor((1, 3, 32, 32)  ,  rng=random.Random(lowercase__						)						).to(lowercase__						)
										_lowerCAmelCase							=							image.cpu().permute(0  ,  2  ,  3  ,  1						)[0]
										_lowerCAmelCase							=							Image.fromarray(np.uinta(lowercase__						)						).convert('RGB'						)
										if str(lowercase__						).startswith('mps'						):
																	_lowerCAmelCase							=							torch.manual_seed(lowercase__						)
										else:
																	_lowerCAmelCase							=							torch.Generator(device=lowercase__						).manual_seed(lowercase__						)
										_lowerCAmelCase							=							{
										    'image': image,
										    'source_prompt': 'a cat and a frog',
										    'target_prompt': 'a dog and a newt',
										    'generator': generator,
										    'num_inference_steps': 2,
										    'num_maps_per_mask': 2,
										    'mask_encode_strength': 1.0,
										    'guidance_scale': 6.0,
										    'output_type': 'numpy',
										}

										return inputs





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	List[str]=0						):
										_lowerCAmelCase							=							floats_tensor((1, 3, 32, 32)  ,  rng=random.Random(lowercase__						)						).to(lowercase__						)
										_lowerCAmelCase							=							image.cpu().permute(0  ,  2  ,  3  ,  1						)[0]
										_lowerCAmelCase							=							Image.fromarray(np.uinta(lowercase__						)						).convert('RGB'						)
										if str(lowercase__						).startswith('mps'						):
																	_lowerCAmelCase							=							torch.manual_seed(lowercase__						)
										else:
																	_lowerCAmelCase							=							torch.Generator(device=lowercase__						).manual_seed(lowercase__						)
										_lowerCAmelCase							=							{
										    'image': image,
										    'prompt': 'a cat and a frog',
										    'generator': generator,
										    'num_inference_steps': 2,
										    'inpaint_strength': 1.0,
										    'guidance_scale': 6.0,
										    'decode_latents': True,
										    'output_type': 'numpy',
										}
										return inputs





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]						):
										if not hasattr(self.pipeline_class  ,  '_optional_components'						):
																	return

										_lowerCAmelCase							=							self.get_dummy_components()
										_lowerCAmelCase							=							self.pipeline_class(**lowercase__						)
										pipe.to(lowercase__						)
										pipe.set_progress_bar_config(disable=lowercase__						)

										# set all optional components to None and update pipeline config accordingly
										for optional_component in pipe._optional_components:
																	setattr(lowercase__  ,  lowercase__  ,  lowercase__						)
										pipe.register_modules(**{optional_component: None for optional_component in pipe._optional_components}						)

										_lowerCAmelCase							=							self.get_dummy_inputs(lowercase__						)
										_lowerCAmelCase							=							pipe(**lowercase__						)[0]

										with tempfile.TemporaryDirectory() as tmpdir:
																	pipe.save_pretrained(lowercase__						)
																	_lowerCAmelCase							=							self.pipeline_class.from_pretrained(lowercase__						)
																	pipe_loaded.to(lowercase__						)
																	pipe_loaded.set_progress_bar_config(disable=lowercase__						)

										for optional_component in pipe._optional_components:
																	self.assertTrue(
																	    getattr(lowercase__  ,  lowercase__						) is None  ,  f'`{optional_component}` did not stay set to None after loading.'  ,  )

										_lowerCAmelCase							=							self.get_dummy_inputs(lowercase__						)
										_lowerCAmelCase							=							pipe_loaded(**lowercase__						)[0]

										_lowerCAmelCase							=							np.abs(output - output_loaded						).max()
										self.assertLess(lowercase__  ,  1e-4						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	str						):
										_lowerCAmelCase							=							'cpu'

										_lowerCAmelCase							=							self.get_dummy_components()
										_lowerCAmelCase							=							self.pipeline_class(**lowercase__						)
										pipe.to(lowercase__						)
										pipe.set_progress_bar_config(disable=lowercase__						)

										_lowerCAmelCase							=							self.get_dummy_mask_inputs(lowercase__						)
										_lowerCAmelCase							=							pipe.generate_mask(**lowercase__						)
										_lowerCAmelCase							=							mask[0, -3:, -3:]

										self.assertEqual(mask.shape  ,  (1, 16, 16)						)
										_lowerCAmelCase							=							np.array([0] * 9						)
										_lowerCAmelCase							=							np.abs(mask_slice.flatten() - expected_slice						).max()
										self.assertLessEqual(lowercase__  ,  1e-3						)
										self.assertEqual(mask[0, -3, -4]  ,  0						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]						):
										_lowerCAmelCase							=							'cpu'

										_lowerCAmelCase							=							self.get_dummy_components()
										_lowerCAmelCase							=							self.pipeline_class(**lowercase__						)
										pipe.to(lowercase__						)
										pipe.set_progress_bar_config(disable=lowercase__						)

										_lowerCAmelCase							=							self.get_dummy_inversion_inputs(lowercase__						)
										_lowerCAmelCase							=							pipe.invert(**lowercase__						).images
										_lowerCAmelCase							=							image[0, -1, -3:, -3:]

										self.assertEqual(image.shape  ,  (2, 32, 32, 3)						)
										_lowerCAmelCase							=							np.array(
										    [0.5_1_5_0, 0.5_1_3_4, 0.5_0_4_3, 0.5_3_7_6, 0.4_6_9_4, 0.5_1_0_5_0, 0.5_0_1_5, 0.4_4_0_7, 0.4_7_9_9]  ,  )
										_lowerCAmelCase							=							np.abs(image_slice.flatten() - expected_slice						).max()
										self.assertLessEqual(lowercase__  ,  1e-3						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	str						):
										super().test_inference_batch_single_identical(expected_max_diff=5e-3						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]						):
										_lowerCAmelCase							=							'cpu'

										_lowerCAmelCase							=							self.get_dummy_components()

										_lowerCAmelCase							=							{'beta_start': 0.0_0_0_8_5, 'beta_end': 0.0_1_2, 'beta_schedule': 'scaled_linear'}
										_lowerCAmelCase							=							DPMSolverMultistepScheduler(**lowercase__						)
										_lowerCAmelCase							=							DPMSolverMultistepInverseScheduler(**lowercase__						)

										_lowerCAmelCase							=							self.pipeline_class(**lowercase__						)
										pipe.to(lowercase__						)
										pipe.set_progress_bar_config(disable=lowercase__						)

										_lowerCAmelCase							=							self.get_dummy_inversion_inputs(lowercase__						)
										_lowerCAmelCase							=							pipe.invert(**lowercase__						).images
										_lowerCAmelCase							=							image[0, -1, -3:, -3:]

										self.assertEqual(image.shape  ,  (2, 32, 32, 3)						)
										_lowerCAmelCase							=							np.array(
										    [0.5_1_5_0, 0.5_1_3_4, 0.5_0_4_3, 0.5_3_7_6, 0.4_6_9_4, 0.5_1_0_5_0, 0.5_0_1_5, 0.4_4_0_7, 0.4_7_9_9]  ,  )
										_lowerCAmelCase							=							np.abs(image_slice.flatten() - expected_slice						).max()
										self.assertLessEqual(lowercase__  ,  1e-3						)







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





			def   SCREAMING_SNAKE_CASE__     (       self	:	Tuple						):
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()





			@classmethod
			def   SCREAMING_SNAKE_CASE__     (       cls	:	int						):
										_lowerCAmelCase							=							load_image(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png'						)

										_lowerCAmelCase							=							raw_image.convert('RGB'						).resize((7_68, 7_68)						)

										_lowerCAmelCase							=							raw_image





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]						):
										_lowerCAmelCase							=							torch.manual_seed(0						)

										_lowerCAmelCase							=							StableDiffusionDiffEditPipeline.from_pretrained(
										    'stabilityai/stable-diffusion-2-1'  ,  safety_checker=lowercase__  ,  torch_dtype=torch.floataa						)
										_lowerCAmelCase							=							DDIMScheduler.from_config(pipe.scheduler.config						)
										_lowerCAmelCase							=							DDIMInverseScheduler.from_config(pipe.scheduler.config						)
										pipe.enable_model_cpu_offload()
										pipe.set_progress_bar_config(disable=lowercase__						)

										_lowerCAmelCase							=							'a bowl of fruit'
										_lowerCAmelCase							=							'a bowl of pears'

										_lowerCAmelCase							=							pipe.generate_mask(
										    image=self.raw_image  ,  source_prompt=lowercase__  ,  target_prompt=lowercase__  ,  generator=lowercase__  ,  )

										_lowerCAmelCase							=							pipe.invert(
										    prompt=lowercase__  ,  image=self.raw_image  ,  inpaint_strength=0.7  ,  generator=lowercase__						).latents

										_lowerCAmelCase							=							pipe(
										    prompt=lowercase__  ,  mask_image=lowercase__  ,  image_latents=lowercase__  ,  generator=lowercase__  ,  negative_prompt=lowercase__  ,  inpaint_strength=0.7  ,  output_type='numpy'  ,  ).images[0]

										_lowerCAmelCase							=							(
										    np.array(
										        load_image(
										            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
										            '/diffedit/pears.png'						).resize((7_68, 7_68)						)						)
										    / 2_55
										)
										assert np.abs((expected_image - image).max()						) < 5e-1





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]						):
										_lowerCAmelCase							=							torch.manual_seed(0						)

										_lowerCAmelCase							=							StableDiffusionDiffEditPipeline.from_pretrained(
										    'stabilityai/stable-diffusion-2-1'  ,  safety_checker=lowercase__  ,  torch_dtype=torch.floataa						)
										_lowerCAmelCase							=							DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
										_lowerCAmelCase							=							DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config						)
										pipe.enable_model_cpu_offload()
										pipe.set_progress_bar_config(disable=lowercase__						)

										_lowerCAmelCase							=							'a bowl of fruit'
										_lowerCAmelCase							=							'a bowl of pears'

										_lowerCAmelCase							=							pipe.generate_mask(
										    image=self.raw_image  ,  source_prompt=lowercase__  ,  target_prompt=lowercase__  ,  generator=lowercase__  ,  )

										_lowerCAmelCase							=							pipe.invert(
										    prompt=lowercase__  ,  image=self.raw_image  ,  inpaint_strength=0.7  ,  generator=lowercase__  ,  num_inference_steps=25  ,  ).latents

										_lowerCAmelCase							=							pipe(
										    prompt=lowercase__  ,  mask_image=lowercase__  ,  image_latents=lowercase__  ,  generator=lowercase__  ,  negative_prompt=lowercase__  ,  inpaint_strength=0.7  ,  num_inference_steps=25  ,  output_type='numpy'  ,  ).images[0]

										_lowerCAmelCase							=							(
										    np.array(
										        load_image(
										            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
										            '/diffedit/pears.png'						).resize((7_68, 7_68)						)						)
										    / 2_55
										)
										assert np.abs((expected_image - image).max()						) < 5e-1







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        lowerCamelCase__				(   UpperCAmelCase				):





			def __init__(       self	:	Dict  ,  lowercase__	:	Dict  ,  lowercase__	:	Optional[Any]=13  ,  lowercase__	:	Dict=7  ,  lowercase__	:	Dict=True  ,  lowercase__	:	Optional[Any]=True  ,  lowercase__	:	Optional[int]=False  ,  lowercase__	:	Any=True  ,  lowercase__	:	Union[str, Any]=99  ,  lowercase__	:	Optional[int]=32  ,  lowercase__	:	Any=5  ,  lowercase__	:	Any=4  ,  lowercase__	:	List[str]=64  ,  lowercase__	:	Any="gelu"  ,  lowercase__	:	Optional[Any]=0.1  ,  lowercase__	:	List[str]=0.1  ,  lowercase__	:	Dict=5_12  ,  lowercase__	:	List[str]=16  ,  lowercase__	:	Union[str, Any]=2  ,  lowercase__	:	str=0.0_2  ,  lowercase__	:	Optional[int]=3  ,  lowercase__	:	Union[str, Any]=4  ,  lowercase__	:	Union[str, Any]=None  ,  lowercase__	:	Optional[int]=2  ,  lowercase__	:	Optional[int]=2  ,  lowercase__	:	List[Any]=2  ,  lowercase__	:	Optional[int]=2  ,  lowercase__	:	Union[str, Any]=4  ,  lowercase__	:	Tuple=1  ,  ):
										_lowerCAmelCase							=							parent
										_lowerCAmelCase							=							batch_size
										_lowerCAmelCase							=							seq_length
										_lowerCAmelCase							=							is_training
										_lowerCAmelCase							=							use_input_mask
										_lowerCAmelCase							=							use_token_type_ids
										_lowerCAmelCase							=							use_labels
										_lowerCAmelCase							=							vocab_size
										_lowerCAmelCase							=							hidden_size
										_lowerCAmelCase							=							num_hidden_layers
										_lowerCAmelCase							=							num_attention_heads
										_lowerCAmelCase							=							intermediate_size
										_lowerCAmelCase							=							hidden_act
										_lowerCAmelCase							=							hidden_dropout_prob
										_lowerCAmelCase							=							attention_probs_dropout_prob
										_lowerCAmelCase							=							max_position_embeddings
										_lowerCAmelCase							=							type_vocab_size
										_lowerCAmelCase							=							type_sequence_label_size
										_lowerCAmelCase							=							initializer_range
										_lowerCAmelCase							=							num_labels
										_lowerCAmelCase							=							num_choices
										_lowerCAmelCase							=							scope
										_lowerCAmelCase							=							q_groups
										_lowerCAmelCase							=							k_groups
										_lowerCAmelCase							=							v_groups
										_lowerCAmelCase							=							post_attention_groups
										_lowerCAmelCase							=							intermediate_groups
										_lowerCAmelCase							=							output_groups





			def   SCREAMING_SNAKE_CASE__     (       self	:	Any						):
										_lowerCAmelCase							=							ids_tensor([self.batch_size, self.seq_length]  ,  self.vocab_size						)

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

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

										_lowerCAmelCase							=							self.get_config()

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





			def   SCREAMING_SNAKE_CASE__     (       self	:	int						):
										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	:	int  ,  lowercase__	:	Dict  ,  lowercase__	:	Any  ,  lowercase__	:	str  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	Any  ,  lowercase__	:	Tuple						):
										_lowerCAmelCase							=							SqueezeBertModel(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(lowercase__  ,  lowercase__						)
										_lowerCAmelCase							=							model(lowercase__						)
										self.parent.assertEqual(result.last_hidden_state.shape  ,  (self.batch_size, self.seq_length, self.hidden_size)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	int  ,  lowercase__	:	List[Any]  ,  lowercase__	:	Dict  ,  lowercase__	:	Union[str, Any]						):
										_lowerCAmelCase							=							SqueezeBertForMaskedLM(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__  ,  labels=lowercase__						)
										self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.vocab_size)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	int  ,  lowercase__	:	Tuple  ,  lowercase__	:	Dict  ,  lowercase__	:	str  ,  lowercase__	:	Dict  ,  lowercase__	:	Tuple  ,  lowercase__	:	int						):
										_lowerCAmelCase							=							SqueezeBertForQuestionAnswering(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(
										    lowercase__  ,  attention_mask=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	:	Dict  ,  lowercase__	:	Dict  ,  lowercase__	:	Tuple  ,  lowercase__	:	Any  ,  lowercase__	:	List[str]  ,  lowercase__	:	List[Any]  ,  lowercase__	:	Tuple						):
										_lowerCAmelCase							=							self.num_labels
										_lowerCAmelCase							=							SqueezeBertForSequenceClassification(lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__  ,  labels=lowercase__						)
										self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	str  ,  lowercase__	:	Dict  ,  lowercase__	:	Tuple  ,  lowercase__	:	Any  ,  lowercase__	:	Optional[int]						):
										_lowerCAmelCase							=							self.num_labels
										_lowerCAmelCase							=							SqueezeBertForTokenClassification(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()

										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__  ,  labels=lowercase__						)
										self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.num_labels)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	Tuple  ,  lowercase__	:	List[Any]  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	Tuple  ,  lowercase__	:	Any						):
										_lowerCAmelCase							=							self.num_choices
										_lowerCAmelCase							=							SqueezeBertForMultipleChoice(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							input_ids.unsqueeze(1						).expand(-1  ,  self.num_choices  ,  -1						).contiguous()
										_lowerCAmelCase							=							input_mask.unsqueeze(1						).expand(-1  ,  self.num_choices  ,  -1						).contiguous()
										_lowerCAmelCase							=							model(
										    lowercase__  ,  attention_mask=lowercase__  ,  labels=lowercase__  ,  )
										self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_choices)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]						):
										_lowerCAmelCase							=							self.prepare_config_and_inputs()
										((_lowerCAmelCase)			,		(_lowerCAmelCase)			,		(_lowerCAmelCase)			,		(_lowerCAmelCase)			,		(_lowerCAmelCase)			,		(_lowerCAmelCase))							=							config_and_inputs
										_lowerCAmelCase							=							{'input_ids': input_ids, 'attention_mask': input_mask}
										return config, inputs_dict







@require_torch
class        lowerCamelCase__				(   UpperCAmelCase       ,UpperCAmelCase       ,unittest.TestCase				):
			UpperCamelCase__    							=(
			    (
			        SqueezeBertModel,
			        SqueezeBertForMaskedLM,
			        SqueezeBertForMultipleChoice,
			        SqueezeBertForQuestionAnswering,
			        SqueezeBertForSequenceClassification,
			        SqueezeBertForTokenClassification,
			    )
			    if is_torch_available()
			    else None
			)
			UpperCamelCase__    							=(
			    {
			        "feature-extraction": SqueezeBertModel,
			        "fill-mask": SqueezeBertForMaskedLM,
			        "question-answering": SqueezeBertForQuestionAnswering,
			        "text-classification": SqueezeBertForSequenceClassification,
			        "token-classification": SqueezeBertForTokenClassification,
			        "zero-shot": SqueezeBertForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			UpperCamelCase__    							=False
			UpperCamelCase__    							=True
			UpperCamelCase__    							=False





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]						):
										_lowerCAmelCase							=							SqueezeBertModelTester(self						)
										_lowerCAmelCase							=							ConfigTester(self  ,  config_class=lowercase__  ,  dim=37						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]						):
										self.config_tester.run_common_tests()





			def   SCREAMING_SNAKE_CASE__     (       self	:	Dict						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_squeezebert_model(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Tuple						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_squeezebert_for_masked_lm(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_squeezebert_for_question_answering(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Any						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_squeezebert_for_sequence_classification(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Any						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_squeezebert_for_token_classification(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Dict						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_squeezebert_for_multiple_choice(*lowercase__						)





			@slow
			def   SCREAMING_SNAKE_CASE__     (       self	:	int						):
										for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	_lowerCAmelCase							=							SqueezeBertModel.from_pretrained(lowercase__						)
																	self.assertIsNotNone(lowercase__						)







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





			@slow
			def   SCREAMING_SNAKE_CASE__     (       self	:	Dict						):
										_lowerCAmelCase							=							SqueezeBertForSequenceClassification.from_pretrained('squeezebert/squeezebert-mnli'						)

										_lowerCAmelCase							=							torch.tensor([[1, 2_94_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 13, 15_88, 2]]						)
										_lowerCAmelCase							=							model(lowercase__						)[0]
										_lowerCAmelCase							=							torch.Size((1, 3)						)
										self.assertEqual(output.shape  ,  lowercase__						)
										_lowerCAmelCase							=							torch.tensor([[0.6_4_0_1, -0.0_3_4_9, -0.6_0_4_1]]						)
										self.assertTrue(torch.allclose(lowercase__  ,  lowercase__  ,  atol=1e-4						)						)






