





import gc
import unittest

import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    ClapTextConfig,
    ClapTextModelWithProjection,
    RobertaTokenizer,
    SpeechTaHifiGan,
    SpeechTaHifiGanConfig,
)

from diffusers import (
    AudioLDMPipeline,
    AutoencoderKL,
    DDIMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class 							__UpperCAmelCase							(						__snake_case    ,  unittest.TestCase   ):



						"""simple docstring"""







						lowercase							  =    AudioLDMPipeline
						lowercase							  =    TEXT_TO_AUDIO_PARAMS
						lowercase							  =    TEXT_TO_AUDIO_BATCH_PARAMS
						lowercase							  =    frozenset(
						    [
						        """num_inference_steps""",
						        """num_waveforms_per_prompt""",
						        """generator""",
						        """latents""",
						        """output_type""",
						        """return_dict""",
						        """callback""",
						        """callback_steps""",
						    ]   )




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





												"""simple docstring"""
												torch.manual_seed(0				)
												UpperCamelCase  							= 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, 64)	,		class_embed_type="simple_projection"	,		projection_class_embeddings_input_dim=32	,		class_embeddings_concat=SCREAMING_SNAKE_CASE	,		)
												UpperCamelCase  							= DDIMScheduler(
												    beta_start=0.00_085	,		beta_end=0.012	,		beta_schedule="scaled_linear"	,		clip_sample=SCREAMING_SNAKE_CASE	,		set_alpha_to_one=SCREAMING_SNAKE_CASE	,		)
												torch.manual_seed(0				)
												UpperCamelCase  							= AutoencoderKL(
												    block_out_channels=[32, 64]	,		in_channels=1	,		out_channels=1	,		down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]	,		up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]	,		latent_channels=4	,		)
												torch.manual_seed(0				)
												UpperCamelCase  							= ClapTextConfig(
												    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=1000	,		projection_dim=32	,		)
												UpperCamelCase  							= ClapTextModelWithProjection(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= RobertaTokenizer.from_pretrained("hf-internal-testing/tiny-random-roberta"	,		model_max_length=77				)

												UpperCamelCase  							= SpeechTaHifiGanConfig(
												    model_in_dim=8	,		sampling_rate=16000	,		upsample_initial_channel=16	,		upsample_rates=[2, 2]	,		upsample_kernel_sizes=[4, 4]	,		resblock_kernel_sizes=[3, 7]	,		resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]]	,		normalize_before=SCREAMING_SNAKE_CASE	,		)

												UpperCamelCase  							= SpeechTaHifiGan(SCREAMING_SNAKE_CASE				)

												UpperCamelCase  							= {
												    "unet": unet,
												    "scheduler": scheduler,
												    "vae": vae,
												    "text_encoder": text_encoder,
												    "tokenizer": tokenizer,
												    "vocoder": vocoder,
												}
												return components




						def 		__lowerCAmelCase			(   self	,		SCREAMING_SNAKE_CASE	,		SCREAMING_SNAKE_CASE=0				) ->      Tuple:





												"""simple docstring"""
												if str(SCREAMING_SNAKE_CASE				).startswith("mps"				):
																		UpperCamelCase  							= torch.manual_seed(SCREAMING_SNAKE_CASE				)
												else:
																		UpperCamelCase  							= torch.Generator(device=SCREAMING_SNAKE_CASE				).manual_seed(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= {
												    "prompt": "A hammer hitting a wooden surface",
												    "generator": generator,
												    "num_inference_steps": 2,
												    "guidance_scale": 6.0,
												}
												return inputs




						def 		__lowerCAmelCase			(   self				) ->      Dict:





												"""simple docstring"""
												UpperCamelCase  							= "cpu"  # ensure determinism for the device-dependent torch.Generator

												UpperCamelCase  							= self.get_dummy_components()
												UpperCamelCase  							= AudioLDMPipeline(**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe.to(SCREAMING_SNAKE_CASE				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE				)

												UpperCamelCase  							= self.get_dummy_inputs(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe(**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= output.audios[0]

												assert audio.ndim == 1
												assert len(SCREAMING_SNAKE_CASE				) == 256

												UpperCamelCase  							= audio[:10]
												UpperCamelCase  							= np.array(
												    [-0.0_050, 0.0_050, -0.0_060, 0.0_033, -0.0_026, 0.0_033, -0.0_027, 0.0_033, -0.0_028, 0.0_033]				)

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




						def 		__lowerCAmelCase			(   self				) ->      str:





												"""simple docstring"""
												UpperCamelCase  							= self.get_dummy_components()
												UpperCamelCase  							= AudioLDMPipeline(**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe.to(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe.to(SCREAMING_SNAKE_CASE				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE				)

												UpperCamelCase  							= self.get_dummy_inputs(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= 3 * [inputs["prompt"]]

												# forward
												UpperCamelCase  							= audioldm_pipe(**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= output.audios[0]

												UpperCamelCase  							= self.get_dummy_inputs(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= 3 * [inputs.pop("prompt"				)]

												UpperCamelCase  							= audioldm_pipe.tokenizer(
												    SCREAMING_SNAKE_CASE	,		padding="max_length"	,		max_length=audioldm_pipe.tokenizer.model_max_length	,		truncation=SCREAMING_SNAKE_CASE	,		return_tensors="pt"	,		)
												UpperCamelCase  							= text_inputs["input_ids"].to(SCREAMING_SNAKE_CASE				)

												UpperCamelCase  							= audioldm_pipe.text_encoder(
												    SCREAMING_SNAKE_CASE	,		)
												UpperCamelCase  							= prompt_embeds.text_embeds
												# additional L_2 normalization over each hidden-state
												UpperCamelCase  							= F.normalize(SCREAMING_SNAKE_CASE	,		dim=-1				)

												UpperCamelCase  							= prompt_embeds

												# forward
												UpperCamelCase  							= audioldm_pipe(**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= output.audios[0]

												assert np.abs(audio_a - audio_a				).max() < 1e-2




						def 		__lowerCAmelCase			(   self				) ->      str:





												"""simple docstring"""
												UpperCamelCase  							= self.get_dummy_components()
												UpperCamelCase  							= AudioLDMPipeline(**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe.to(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe.to(SCREAMING_SNAKE_CASE				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE				)

												UpperCamelCase  							= self.get_dummy_inputs(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= 3 * ["this is a negative prompt"]
												UpperCamelCase  							= negative_prompt
												UpperCamelCase  							= 3 * [inputs["prompt"]]

												# forward
												UpperCamelCase  							= audioldm_pipe(**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= output.audios[0]

												UpperCamelCase  							= self.get_dummy_inputs(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= 3 * [inputs.pop("prompt"				)]

												UpperCamelCase  							= []
												for p in [prompt, negative_prompt]:
																		UpperCamelCase  							= audioldm_pipe.tokenizer(
																		    SCREAMING_SNAKE_CASE	,		padding="max_length"	,		max_length=audioldm_pipe.tokenizer.model_max_length	,		truncation=SCREAMING_SNAKE_CASE	,		return_tensors="pt"	,		)
																		UpperCamelCase  							= text_inputs["input_ids"].to(SCREAMING_SNAKE_CASE				)

																		UpperCamelCase  							= audioldm_pipe.text_encoder(
																		    SCREAMING_SNAKE_CASE	,		)
																		UpperCamelCase  							= text_embeds.text_embeds
																		# additional L_2 normalization over each hidden-state
																		UpperCamelCase  							= F.normalize(SCREAMING_SNAKE_CASE	,		dim=-1				)

																		embeds.append(SCREAMING_SNAKE_CASE				)

												UpperCamelCase  							= embeds

												# forward
												UpperCamelCase  							= audioldm_pipe(**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= output.audios[0]

												assert np.abs(audio_a - audio_a				).max() < 1e-2




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





												"""simple docstring"""
												UpperCamelCase  							= "cpu"  # ensure determinism for the device-dependent torch.Generator
												UpperCamelCase  							= self.get_dummy_components()
												UpperCamelCase  							= PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= AudioLDMPipeline(**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe.to(SCREAMING_SNAKE_CASE				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE				)

												UpperCamelCase  							= self.get_dummy_inputs(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= "egg cracking"
												UpperCamelCase  							= audioldm_pipe(**SCREAMING_SNAKE_CASE	,		negative_prompt=SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= output.audios[0]

												assert audio.ndim == 1
												assert len(SCREAMING_SNAKE_CASE				) == 256

												UpperCamelCase  							= audio[:10]
												UpperCamelCase  							= np.array(
												    [-0.0_051, 0.0_050, -0.0_060, 0.0_034, -0.0_026, 0.0_033, -0.0_027, 0.0_033, -0.0_028, 0.0_032]				)

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




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





												"""simple docstring"""
												UpperCamelCase  							= "cpu"  # ensure determinism for the device-dependent torch.Generator
												UpperCamelCase  							= self.get_dummy_components()
												UpperCamelCase  							= PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= AudioLDMPipeline(**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe.to(SCREAMING_SNAKE_CASE				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE				)

												UpperCamelCase  							= "A hammer hitting a wooden surface"

												# test num_waveforms_per_prompt=1 (default)
												UpperCamelCase  							= audioldm_pipe(SCREAMING_SNAKE_CASE	,		num_inference_steps=2				).audios

												assert audios.shape == (1, 256)

												# test num_waveforms_per_prompt=1 (default) for batch of prompts
												UpperCamelCase  							= 2
												UpperCamelCase  							= audioldm_pipe([prompt] * batch_size	,		num_inference_steps=2				).audios

												assert audios.shape == (batch_size, 256)

												# test num_waveforms_per_prompt for single prompt
												UpperCamelCase  							= 2
												UpperCamelCase  							= audioldm_pipe(SCREAMING_SNAKE_CASE	,		num_inference_steps=2	,		num_waveforms_per_prompt=SCREAMING_SNAKE_CASE				).audios

												assert audios.shape == (num_waveforms_per_prompt, 256)

												# test num_waveforms_per_prompt for batch of prompts
												UpperCamelCase  							= 2
												UpperCamelCase  							= audioldm_pipe(
												    [prompt] * batch_size	,		num_inference_steps=2	,		num_waveforms_per_prompt=SCREAMING_SNAKE_CASE				).audios

												assert audios.shape == (batch_size * num_waveforms_per_prompt, 256)




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





												"""simple docstring"""
												UpperCamelCase  							= "cpu"  # ensure determinism for the device-dependent torch.Generator
												UpperCamelCase  							= self.get_dummy_components()
												UpperCamelCase  							= AudioLDMPipeline(**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe.to(SCREAMING_SNAKE_CASE				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe.vocoder.config.sampling_rate

												UpperCamelCase  							= self.get_dummy_inputs(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe(audio_length_in_s=0.016	,		**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= output.audios[0]

												assert audio.ndim == 1
												assert len(SCREAMING_SNAKE_CASE				) / vocoder_sampling_rate == 0.016

												UpperCamelCase  							= audioldm_pipe(audio_length_in_s=0.032	,		**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= output.audios[0]

												assert audio.ndim == 1
												assert len(SCREAMING_SNAKE_CASE				) / vocoder_sampling_rate == 0.032




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





												"""simple docstring"""
												UpperCamelCase  							= self.get_dummy_components()
												UpperCamelCase  							= AudioLDMPipeline(**SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe.to(SCREAMING_SNAKE_CASE				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE				)

												UpperCamelCase  							= ["hey"]

												UpperCamelCase  							= audioldm_pipe(SCREAMING_SNAKE_CASE	,		num_inference_steps=1				)
												UpperCamelCase  							= output.audios.shape
												assert audio_shape == (1, 256)

												UpperCamelCase  							= audioldm_pipe.vocoder.config
												config.model_in_dim *= 2
												UpperCamelCase  							= SpeechTaHifiGan(SCREAMING_SNAKE_CASE				).to(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe(SCREAMING_SNAKE_CASE	,		num_inference_steps=1				)
												UpperCamelCase  							= output.audios.shape
												# waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
												assert audio_shape == (1, 256)




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





												"""simple docstring"""
												self._test_attention_slicing_forward_pass(test_mean_pixel_difference=SCREAMING_SNAKE_CASE				)




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





												"""simple docstring"""
												self._test_inference_batch_single_identical(test_mean_pixel_difference=SCREAMING_SNAKE_CASE				)




						@unittest.skipIf(
						    torch_device != "cuda" or not is_xformers_available()	,		reason="XFormers attention is only available with CUDA and `xformers` installed"	,		)
						def 		__lowerCAmelCase			(   self				) ->      List[Any]:





												"""simple docstring"""
												self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=SCREAMING_SNAKE_CASE				)


@slow
class 							__UpperCAmelCase							(						unittest.TestCase   ):



						"""simple docstring"""







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





												"""simple docstring"""
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()




						def 		__lowerCAmelCase			(   self	,		SCREAMING_SNAKE_CASE	,		SCREAMING_SNAKE_CASE="cpu"	,		SCREAMING_SNAKE_CASE=torch.floataa	,		SCREAMING_SNAKE_CASE=0				) ->      Dict:





												"""simple docstring"""
												UpperCamelCase  							= torch.Generator(device=SCREAMING_SNAKE_CASE				).manual_seed(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= np.random.RandomState(SCREAMING_SNAKE_CASE				).standard_normal((1, 8, 128, 16)				)
												UpperCamelCase  							= torch.from_numpy(SCREAMING_SNAKE_CASE				).to(device=SCREAMING_SNAKE_CASE	,		dtype=SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= {
												    "prompt": "A hammer hitting a wooden surface",
												    "latents": latents,
												    "generator": generator,
												    "num_inference_steps": 3,
												    "guidance_scale": 2.5,
												}
												return inputs




						def 		__lowerCAmelCase			(   self				) ->      Any:





												"""simple docstring"""
												UpperCamelCase  							= AudioLDMPipeline.from_pretrained("cvssp/audioldm"				)
												UpperCamelCase  							= audioldm_pipe.to(SCREAMING_SNAKE_CASE				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE				)

												UpperCamelCase  							= self.get_inputs(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= 25
												UpperCamelCase  							= audioldm_pipe(**SCREAMING_SNAKE_CASE				).audios[0]

												assert audio.ndim == 1
												assert len(SCREAMING_SNAKE_CASE				) == 81920

												UpperCamelCase  							= audio[77230:77240]
												UpperCamelCase  							= np.array(
												    [-0.4_884, -0.4_607, 0.0_023, 0.5_007, 0.5_896, 0.5_151, 0.3_813, -0.0_208, -0.3_687, -0.4_315]				)
												UpperCamelCase  							= np.abs(expected_slice - audio_slice				).max()
												assert max_diff < 1e-2




						def 		__lowerCAmelCase			(   self				) ->      Dict:





												"""simple docstring"""
												UpperCamelCase  							= AudioLDMPipeline.from_pretrained("cvssp/audioldm"				)
												UpperCamelCase  							= LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config				)
												UpperCamelCase  							= audioldm_pipe.to(SCREAMING_SNAKE_CASE				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE				)

												UpperCamelCase  							= self.get_inputs(SCREAMING_SNAKE_CASE				)
												UpperCamelCase  							= audioldm_pipe(**SCREAMING_SNAKE_CASE				).audios[0]

												assert audio.ndim == 1
												assert len(SCREAMING_SNAKE_CASE				) == 81920

												UpperCamelCase  							= audio[27780:27790]
												UpperCamelCase  							= np.array([-0.2_131, -0.0_873, -0.0_124, -0.0_189, 0.0_569, 0.1_373, 0.1_883, 0.2_886, 0.3_297, 0.2_212]				)
												UpperCamelCase  							= np.abs(expected_slice - audio_slice				).max()
												assert max_diff < 3e-2






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   __lowercase						(						__snake_case       ):






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






	def       _a(self				: int    ) ->			Dict:
							_lowercase       :	str           =   ids_tensor([self.batch_size, self.seq_length]      ,	self.vocab_size    )

							_lowercase       :	Tuple           =   None
							if self.use_input_mask:
													_lowercase       :	int           =   random_attention_mask([self.batch_size, self.seq_length]    )

							_lowercase       :	Tuple           =   None
							_lowercase       :	Union[str, Any]           =   None
							_lowercase       :	Tuple           =   None
							if self.use_labels:
													_lowercase       :	Dict           =   ids_tensor([self.batch_size]      ,	self.type_sequence_label_size    )
													_lowercase       :	Any           =   ids_tensor([self.batch_size, self.seq_length]      ,	self.num_labels    )
													_lowercase       :	str           =   ids_tensor([self.batch_size]      ,	self.num_choices    )

							_lowercase       :	Dict           =   self.get_config()

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






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






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






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






	def       _a(self				: int      ,	snake_case				: Union[str, Any]      ,	snake_case				: Optional[Any]      ,	snake_case				: int      ,	snake_case				: Optional[int]      ,	snake_case				: Optional[int]      ,	snake_case				: str    ) ->			str:
							_lowercase       :	str           =   self.num_labels
							_lowercase       :	List[str]           =   DistilBertForTokenClassification(config=snake_case    )
							model.to(snake_case    )
							model.eval()

							_lowercase       :	str           =   model(snake_case      ,	attention_mask=snake_case      ,	labels=snake_case    )
							self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.seq_length, self.num_labels)    )






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






	def       _a(self				: List[str]    ) ->			List[str]:
							_lowercase       :	Union[str, Any]           =   self.prepare_config_and_inputs()
							((_lowercase)      ,    (_lowercase)      ,    (_lowercase)      ,    (_lowercase)      ,    (_lowercase)      ,    (_lowercase))       :	Union[str, Any]           =   config_and_inputs
							_lowercase       :	Union[str, Any]           =   {"input_ids": input_ids, "attention_mask": input_mask}
							return config, inputs_dict




@require_torch
class   __lowercase						(						__snake_case   , __snake_case   , unittest.TestCase       ):
	_A       =   (
	    (
	        DistilBertModel,
	        DistilBertForMaskedLM,
	        DistilBertForMultipleChoice,
	        DistilBertForQuestionAnswering,
	        DistilBertForSequenceClassification,
	        DistilBertForTokenClassification,
	    )
	    if is_torch_available()
	    else None
	)
	_A       =   (
	    {
	        "feature-extraction": DistilBertModel,
	        "fill-mask": DistilBertForMaskedLM,
	        "question-answering": DistilBertForQuestionAnswering,
	        "text-classification": DistilBertForSequenceClassification,
	        "token-classification": DistilBertForTokenClassification,
	        "zero-shot": DistilBertForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	_A       =   True
	_A       =   True
	_A       =   True
	_A       =   True






	def       _a(self				: Dict    ) ->			List[Any]:
							_lowercase       :	Optional[Any]           =   DistilBertModelTester(self    )
							_lowercase       :	str           =   ConfigTester(self      ,	config_class=snake_case      ,	dim=37    )






	def       _a(self				: int    ) ->			List[str]:
							self.config_tester.run_common_tests()






	def       _a(self				: Optional[Any]    ) ->			Optional[int]:
							_lowercase       :	Tuple           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_model(*snake_case    )






	def       _a(self				: Any    ) ->			int:
							_lowercase       :	str           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_masked_lm(*snake_case    )






	def       _a(self				: Dict    ) ->			List[Any]:
							_lowercase       :	List[str]           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_question_answering(*snake_case    )






	def       _a(self				: str    ) ->			Tuple:
							_lowercase       :	Tuple           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_sequence_classification(*snake_case    )






	def       _a(self				: Any    ) ->			List[Any]:
							_lowercase       :	Optional[Any]           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_token_classification(*snake_case    )






	def       _a(self				: Optional[int]    ) ->			Optional[int]:
							_lowercase       :	Any           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_multiple_choice(*snake_case    )






	@slow
	def       _a(self				: Optional[Any]    ) ->			Dict:
							for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													_lowercase       :	Tuple           =   DistilBertModel.from_pretrained(snake_case    )
													self.assertIsNotNone(snake_case    )






	@slow
	@require_torch_gpu
	def       _a(self				: Optional[int]    ) ->			Optional[int]:
							_lowercase      ,    _lowercase       :	Optional[int]           =   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

													_lowercase       :	str           =   True
													_lowercase       :	Tuple           =   model_class(config=snake_case    )

													_lowercase       :	str           =   self._prepare_for_class(snake_case      ,	snake_case    )
													_lowercase       :	Optional[int]           =   torch.jit.trace(
													    snake_case      ,	(inputs_dict["input_ids"].to("cpu"    ), inputs_dict["attention_mask"].to("cpu"    ))    )

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




@require_torch
class   __lowercase						(						unittest.TestCase       ):






	@slow
	def       _a(self				: int    ) ->			str:
							_lowercase       :	Any           =   DistilBertModel.from_pretrained("distilbert-base-uncased"    )
							_lowercase       :	Dict           =   torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]    )
							_lowercase       :	List[Any]           =   torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]    )
							with torch.no_grad():
													_lowercase       :	Optional[int]           =   model(snake_case      ,	attention_mask=snake_case    )[0]
							_lowercase       :	List[str]           =   torch.Size((1, 11, 768)    )
							self.assertEqual(output.shape      ,	snake_case    )
							_lowercase       :	Any           =   torch.tensor(
							    [[[-0.16_39, 0.32_99, 0.16_48], [-0.17_46, 0.32_89, 0.17_10], [-0.18_84, 0.33_57, 0.18_10]]]    )

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




