
import inspect
from typing import Callable, List, Optional, Union

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import DiffusionPipeline
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import logging


SCREAMING_SNAKE_CASE				:Optional[int]									=							logging.get_logger(__name__)  # pylint: disable=invalid-name





class 				UpperCAmelCase (    __SCREAMING_SNAKE_CASE    ):




	'''simple docstring'''





	def __init__(			self   :					Optional[Any]				,A   :					AutoencoderKL				,A   :					CLIPTextModel				,A   :					CLIPTokenizer				,A   :					UNetaDConditionModel				,A   :					Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]				,A   :					StableDiffusionSafetyChecker				,A   :					CLIPImageProcessor				,):
					super().__init__()
					self.register_modules(
					    vae=A				,text_encoder=A				,tokenizer=A				,unet=A				,scheduler=A				,safety_checker=A				,feature_extractor=A				,)






	def 							UpperCamelCase_    (			self   :					Tuple				,A   :					Optional[Union[str, int]] = "auto"					):
					if slice_size == "auto":
									# half the attention head size is usually a good trade-off between
									# speed and memory
									__A       =					self.unet.config.attention_head_dim // 2
					self.unet.set_attention_slice(A					)






	def 							UpperCamelCase_    (			self   :					Tuple					):
					self.enable_attention_slicing(A					)







	@torch.no_grad()
	def __call__(			self   :					Dict				,A   :					Union[str, List[str]]				,A   :					int = 5_12				,A   :					int = 5_12				,A   :					int = 50				,A   :					float = 7.5				,A   :					Optional[Union[str, List[str]]] = None				,A   :					Optional[int] = 1				,A   :					float = 0.0				,A   :					Optional[torch.Generator] = None				,A   :					Optional[torch.FloatTensor] = None				,A   :					Optional[str] = "pil"				,A   :					bool = True				,A   :					Optional[Callable[[int, int, torch.FloatTensor], None]] = None				,A   :					int = 1				,A   :					Optional[torch.FloatTensor] = None				,**A   :					Optional[Any]				,):

					if isinstance(A				,A					):
									__A       =					1
					elif isinstance(A				,A					):
									__A       =					len(A					)
					else:
									raise ValueError(f'''`prompt` has to be of type `str` or `list` but is {type(A					)}'''					)

					if height % 8 != 0 or width % 8 != 0:
									raise ValueError(f'''`height` and `width` have to be divisible by 8 but are {height} and {width}.'''					)

					if (callback_steps is None) or (
					    callback_steps is not None and (not isinstance(A				,A					) or callback_steps <= 0)
					):
									raise ValueError(
									    f'''`callback_steps` has to be a positive integer but is {callback_steps} of type'''
									    f''' {type(A					)}.'''					)

					# get prompt text embeddings
					__A       =					self.tokenizer(
					    A				,padding="max_length"				,max_length=self.tokenizer.model_max_length				,return_tensors="pt"				,)
					__A       =					text_inputs.input_ids

					if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
									__A       =					self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]					)
									logger.warning(
									    "The following part of your input was truncated because CLIP can only handle sequences up to"
									    f''' {self.tokenizer.model_max_length} tokens: {removed_text}'''					)
									__A       =					text_input_ids[:, : self.tokenizer.model_max_length]

					if text_embeddings is None:
									__A       =					self.text_encoder(text_input_ids.to(self.device					)					)[0]

					# duplicate text embeddings for each generation per prompt, using mps friendly method
					__A     ,				__A     ,				__A       =					text_embeddings.shape
					__A       =					text_embeddings.repeat(1				,A				,1					)
					__A       =					text_embeddings.view(bs_embed * num_images_per_prompt				,A				,-1					)

					# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
					# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
					# corresponds to doing no classifier free guidance.
					__A       =					guidance_scale > 1.0
					# get unconditional embeddings for classifier free guidance
					if do_classifier_free_guidance:
									__A       =					42
									if negative_prompt is None:
													__A       =					[""]
									elif type(A					) is not type(A					):
													raise TypeError(
													    f'''`negative_prompt` should be the same type to `prompt`, but got {type(A					)} !='''
													    f''' {type(A					)}.'''					)
									elif isinstance(A				,A					):
													__A       =					[negative_prompt]
									elif batch_size != len(A					):
													raise ValueError(
													    f'''`negative_prompt`: {negative_prompt} has batch size {len(A					)}, but `prompt`:'''
													    f''' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'''
													    " the batch size of `prompt`."					)
									else:
													__A       =					negative_prompt

									__A       =					text_input_ids.shape[-1]
									__A       =					self.tokenizer(
									    A				,padding="max_length"				,max_length=A				,truncation=A				,return_tensors="pt"				,)
									__A       =					self.text_encoder(uncond_input.input_ids.to(self.device					)					)[0]

									# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
									__A       =					uncond_embeddings.shape[1]
									__A       =					uncond_embeddings.repeat(A				,A				,1					)
									__A       =					uncond_embeddings.view(batch_size * num_images_per_prompt				,A				,-1					)

									# For classifier free guidance, we need to do two forward passes.
									# Here we concatenate the unconditional and text embeddings into a single batch
									# to avoid doing two forward passes
									__A       =					torch.cat([uncond_embeddings, text_embeddings]					)

					# get the initial random noise unless the user supplied it

					# Unlike in other pipelines, latents need to be generated in the target device
					# for 1-to-1 results reproducibility with the CompVis implementation.
					# However this currently doesn't work in `mps`.
					__A       =					(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
					__A       =					(batch_size * num_images_per_prompt, self.unet.config.in_channels, 64, 64)
					__A       =					text_embeddings.dtype
					if latents is None:
									if self.device.type == "mps":
													# randn does not exist on mps
													__A       =					torch.randn(
													    A				,generator=A				,device="cpu"				,dtype=A					).to(self.device					)
													__A       =					torch.randn(A				,generator=A				,device="cpu"				,dtype=A					).to(
													    self.device					)
									else:
													__A       =					torch.randn(
													    A				,generator=A				,device=self.device				,dtype=A					)
													__A       =					torch.randn(A				,generator=A				,device=self.device				,dtype=A					)
					else:
									if latents_reference.shape != latents_shape:
													raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}'''					)
									__A       =					latents_reference.to(self.device					)
									__A       =					latents.to(self.device					)

					# This is the key part of the pipeline where we
					# try to ensure that the generated images w/ the same seed
					# but different sizes actually result in similar images
					__A       =					(latents_shape[3] - latents_shape_reference[3]) // 2
					__A       =					(latents_shape[2] - latents_shape_reference[2]) // 2
					__A       =					latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
					__A       =					latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
					__A       =					0 if dx < 0 else dx
					__A       =					0 if dy < 0 else dy
					__A       =					max(-dx				,0					)
					__A       =					max(-dy				,0					)
					# import pdb
					# pdb.set_trace()
					__A       =					latents_reference[:, :, dy : dy + h, dx : dx + w]

					# set timesteps
					self.scheduler.set_timesteps(A					)

					# Some schedulers like PNDM have timesteps as arrays
					# It's more optimized to move all timesteps to correct device beforehand
					__A       =					self.scheduler.timesteps.to(self.device					)

					# scale the initial noise by the standard deviation required by the scheduler
					__A       =					latents * self.scheduler.init_noise_sigma

					# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
					# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
					# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
					# and should be between [0, 1]
					__A       =					"eta" in set(inspect.signature(self.scheduler.step					).parameters.keys()					)
					__A       =					{}
					if accepts_eta:
									__A       =					eta

					for i, t in enumerate(self.progress_bar(A					)					):
									# expand the latents if we are doing classifier free guidance
									__A       =					torch.cat([latents] * 2					) if do_classifier_free_guidance else latents
									__A       =					self.scheduler.scale_model_input(A				,A					)

									# predict the noise residual
									__A       =					self.unet(A				,A				,encoder_hidden_states=A					).sample

									# perform guidance
									if do_classifier_free_guidance:
													__A     ,				__A       =					noise_pred.chunk(2					)
													__A       =					noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

									# compute the previous noisy sample x_t -> x_t-1
									__A       =					self.scheduler.step(A				,A				,A				,**A					).prev_sample

									# call the callback, if provided
									if callback is not None and i % callback_steps == 0:
													callback(A				,A				,A					)

					__A       =					1 / 0.1_82_15 * latents
					__A       =					self.vae.decode(A					).sample

					__A       =					(image / 2 + 0.5).clamp(0				,1					)

					# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
					__A       =					image.cpu().permute(0				,2				,3				,1					).float().numpy()

					if self.safety_checker is not None:
									__A       =					self.feature_extractor(self.numpy_to_pil(A					)				,return_tensors="pt"					).to(
									    self.device					)
									__A     ,				__A       =					self.safety_checker(
									    images=A				,clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype					)					)
					else:
									__A       =					None

					if output_type == "pil":
									__A       =					self.numpy_to_pil(A					)

					if not return_dict:
									return (image, has_nsfw_concept)

					return StableDiffusionPipelineOutput(images=A				,nsfw_content_detected=A					)



import os
import tempfile
import unittest

from transformers import DistilBertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, 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 (
			    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
			    DistilBertForMaskedLM,
			    DistilBertForMultipleChoice,
			    DistilBertForQuestionAnswering,
			    DistilBertForSequenceClassification,
			    DistilBertForTokenClassification,
			    DistilBertModel,
			)






class _UpperCamelCase				(				__A ):
			'''simple docstring'''







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

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							parent
				SCREAMING_SNAKE_CASE :   Any					  =							batch_size
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							seq_length
				SCREAMING_SNAKE_CASE :   List[Any]					  =							is_training
				SCREAMING_SNAKE_CASE :   int					  =							use_input_mask
				SCREAMING_SNAKE_CASE :   Tuple					  =							use_token_type_ids
				SCREAMING_SNAKE_CASE :   str					  =							use_labels
				SCREAMING_SNAKE_CASE :   Any					  =							vocab_size
				SCREAMING_SNAKE_CASE :   List[Any]					  =							hidden_size
				SCREAMING_SNAKE_CASE :   str					  =							num_hidden_layers
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							num_attention_heads
				SCREAMING_SNAKE_CASE :   Tuple					  =							intermediate_size
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							hidden_act
				SCREAMING_SNAKE_CASE :   Dict					  =							hidden_dropout_prob
				SCREAMING_SNAKE_CASE :   str					  =							attention_probs_dropout_prob
				SCREAMING_SNAKE_CASE :   Tuple					  =							max_position_embeddings
				SCREAMING_SNAKE_CASE :   List[str]					  =							type_vocab_size
				SCREAMING_SNAKE_CASE :   List[str]					  =							type_sequence_label_size
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							initializer_range
				SCREAMING_SNAKE_CASE :   Tuple					  =							num_labels
				SCREAMING_SNAKE_CASE :   Tuple					  =							num_choices
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							scope
			def 				__UpperCamelCase			(      self				:			Union[str, Any]						)				->							List[str]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							ids_tensor([self.batch_size, self.seq_length]   ,   self.vocab_size						)

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

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

				SCREAMING_SNAKE_CASE :   int					  =							self.get_config()

				return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
			def 				__UpperCamelCase			(      self				:			Dict						)				->							str:

				"""simple docstring"""

				return DistilBertConfig(
				    vocab_size=self.vocab_size   ,   dim=self.hidden_size   ,   n_layers=self.num_hidden_layers   ,   n_heads=self.num_attention_heads   ,   hidden_dim=self.intermediate_size   ,   hidden_act=self.hidden_act   ,   dropout=self.hidden_dropout_prob   ,   attention_dropout=self.attention_probs_dropout_prob   ,   max_position_embeddings=self.max_position_embeddings   ,   initializer_range=self.initializer_range   ,   )
			def 				__UpperCamelCase			(      self				:			Optional[Any]   ,   a				:			int   ,   a				:			Optional[int]   ,   a				:			Optional[int]   ,   a				:			Dict   ,   a				:			str   ,   a				:			str						)				->							int:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   int					  =							DistilBertModel(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							model(a   ,   a						)
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							model(a						)
				self.parent.assertEqual(result.last_hidden_state.shape   ,   (self.batch_size, self.seq_length, self.hidden_size)						)
			def 				__UpperCamelCase			(      self				:			Tuple   ,   a				:			Optional[int]   ,   a				:			Dict   ,   a				:			Tuple   ,   a				:			int   ,   a				:			int   ,   a				:			Any						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							DistilBertForMaskedLM(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   str					  =							model(a   ,   attention_mask=a   ,   labels=a						)
				self.parent.assertEqual(result.logits.shape   ,   (self.batch_size, self.seq_length, self.vocab_size)						)
			def 				__UpperCamelCase			(      self				:			List[Any]   ,   a				:			int   ,   a				:			Optional[Any]   ,   a				:			Optional[Any]   ,   a				:			str   ,   a				:			str   ,   a				:			Tuple						)				->							List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							DistilBertForQuestionAnswering(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   List[str]					  =							model(
				    a   ,   attention_mask=a   ,   start_positions=a   ,   end_positions=a						)
				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 				__UpperCamelCase			(      self				:			Optional[int]   ,   a				:			str   ,   a				:			Any   ,   a				:			int   ,   a				:			Optional[Any]   ,   a				:			int   ,   a				:			str						)				->							Optional[int]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							self.num_labels
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							DistilBertForSequenceClassification(a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							model(a   ,   attention_mask=a   ,   labels=a						)
				self.parent.assertEqual(result.logits.shape   ,   (self.batch_size, self.num_labels)						)
			def 				__UpperCamelCase			(      self				:			Optional[Any]   ,   a				:			List[Any]   ,   a				:			Optional[int]   ,   a				:			Union[str, Any]   ,   a				:			Dict   ,   a				:			Any   ,   a				:			Optional[Any]						)				->							Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   str					  =							self.num_labels
				SCREAMING_SNAKE_CASE :   List[str]					  =							DistilBertForTokenClassification(config=a						)
				model.to(a						)
				model.eval()

				SCREAMING_SNAKE_CASE :   List[str]					  =							model(a   ,   attention_mask=a   ,   labels=a						)
				self.parent.assertEqual(result.logits.shape   ,   (self.batch_size, self.seq_length, self.num_labels)						)
			def 				__UpperCamelCase			(      self				:			int   ,   a				:			Any   ,   a				:			Optional[int]   ,   a				:			Union[str, Any]   ,   a				:			Tuple   ,   a				:			Optional[int]   ,   a				:			Tuple						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[str]					  =							self.num_choices
				SCREAMING_SNAKE_CASE :   Any					  =							DistilBertForMultipleChoice(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							input_ids.unsqueeze(1						).expand(-1   ,   self.num_choices   ,   -1						).contiguous()
				SCREAMING_SNAKE_CASE :   Dict					  =							input_mask.unsqueeze(1						).expand(-1   ,   self.num_choices   ,   -1						).contiguous()
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							model(
				    a   ,   attention_mask=a   ,   labels=a   ,   )
				self.parent.assertEqual(result.logits.shape   ,   (self.batch_size, self.num_choices)						)






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

				"""simple docstring"""

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






@require_torch
class _UpperCamelCase				(				__A				,						__A				,						unittest.TestCase ):
			'''simple docstring'''







			lowerCamelCase__							     =(
			    (
			        DistilBertModel,
			        DistilBertForMaskedLM,
			        DistilBertForMultipleChoice,
			        DistilBertForQuestionAnswering,
			        DistilBertForSequenceClassification,
			        DistilBertForTokenClassification,
			    )
			    if is_torch_available()
			    else None
			)
			lowerCamelCase__							     =(
			    {
			        'feature-extraction': DistilBertModel,
			        'fill-mask': DistilBertForMaskedLM,
			        'question-answering': DistilBertForQuestionAnswering,
			        'text-classification': DistilBertForSequenceClassification,
			        'token-classification': DistilBertForTokenClassification,
			        'zero-shot': DistilBertForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			lowerCamelCase__							     =True
			lowerCamelCase__							     =True
			lowerCamelCase__							     =True
			lowerCamelCase__							     =True
			def 				__UpperCamelCase			(      self				:			Optional[Any]						)				->							Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							DistilBertModelTester(self						)
				SCREAMING_SNAKE_CASE :   List[str]					  =							ConfigTester(self   ,   config_class=a   ,   dim=37						)
			def 				__UpperCamelCase			(      self				:			List[Any]						)				->							Tuple:

				"""simple docstring"""

				self.config_tester.run_common_tests()
			def 				__UpperCamelCase			(      self				:			Union[str, Any]						)				->							List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   int					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_model(*a						)
			def 				__UpperCamelCase			(      self				:			Tuple						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_masked_lm(*a						)
			def 				__UpperCamelCase			(      self				:			List[Any]						)				->							Optional[int]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[Any]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_question_answering(*a						)
			def 				__UpperCamelCase			(      self				:			Union[str, Any]						)				->							Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_sequence_classification(*a						)
			def 				__UpperCamelCase			(      self				:			str						)				->							List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[str]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_token_classification(*a						)
			def 				__UpperCamelCase			(      self				:			List[Any]						)				->							str:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[Any]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_multiple_choice(*a						)
			@slow
			def 				__UpperCamelCase			(      self				:			int						)				->							Any:

				"""simple docstring"""

				for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					SCREAMING_SNAKE_CASE :   Optional[Any]					  =							DistilBertModel.from_pretrained(a						)
					self.assertIsNotNone(a						)






			@slow
			@require_torch_gpu
			def 				__UpperCamelCase			(      self				:			List[Any]						)				->							Optional[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE :   Dict					  =							self.model_tester.prepare_config_and_inputs_for_common()
				for model_class in self.all_model_classes:
					# BertForMultipleChoice behaves incorrectly in JIT environments.
					if model_class == DistilBertForMultipleChoice:
						return

					SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							True
					SCREAMING_SNAKE_CASE :   Any					  =							model_class(config=a						)

					SCREAMING_SNAKE_CASE :   List[Any]					  =							self._prepare_for_class(a   ,   a						)
					SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							torch.jit.trace(
					    a   ,   (inputs_dict["input_ids"].to("cpu"						), inputs_dict["attention_mask"].to("cpu"						))						)

					with tempfile.TemporaryDirectory() as tmp:
						torch.jit.save(a   ,   os.path.join(a   ,   "traced_model.pt"						)						)
						SCREAMING_SNAKE_CASE :   Tuple					  =							torch.jit.load(os.path.join(a   ,   "traced_model.pt"						)   ,   map_location=a						)
						loaded(inputs_dict["input_ids"].to(a						)   ,   inputs_dict["attention_mask"].to(a						)						)






@require_torch
class _UpperCamelCase				(				unittest.TestCase ):
			'''simple docstring'''







			@slow
			def 				__UpperCamelCase			(      self				:			int						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							DistilBertModel.from_pretrained("distilbert-base-uncased"						)
				SCREAMING_SNAKE_CASE :   List[str]					  =							torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]						)
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]						)
				with torch.no_grad():
					SCREAMING_SNAKE_CASE :   Optional[Any]					  =							model(a   ,   attention_mask=a						)[0]
				SCREAMING_SNAKE_CASE :   List[str]					  =							torch.Size((1, 11, 768)						)
				self.assertEqual(output.shape   ,   a						)
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							torch.tensor(
				    [[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]						)

				self.assertTrue(torch.allclose(output[:, 1:4, 1:4]   ,   a   ,   atol=1e-4						)						)