




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 				lowerCAmelCase_       (					__A							):






				def __init__(			self      ,snake_case__      ,snake_case__=13      ,snake_case__=7      ,snake_case__=True      ,snake_case__=True      ,snake_case__=False      ,snake_case__=True      ,snake_case__=99      ,snake_case__=32      ,snake_case__=5      ,snake_case__=4      ,snake_case__=37      ,snake_case__="gelu"      ,snake_case__=0.1      ,snake_case__=0.1      ,snake_case__=512      ,snake_case__=16      ,snake_case__=2      ,snake_case__=0.02      ,snake_case__=3      ,snake_case__=4      ,snake_case__=None      ,):
							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 		snake_case	(			self ):
							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 		snake_case	(			self ):
							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 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : int					  =      DistilBertModel(config=snake_case__ )
							model.to(snake_case__ )
							model.eval()
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      model(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      model(snake_case__ )
							self.parent.assertEqual(result.last_hidden_state.shape      ,(self.batch_size, self.seq_length, self.hidden_size) )






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






				def 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : str					  =      self.num_labels
							SCREAMING_SNAKE_CASE_   : List[str]					  =      DistilBertForTokenClassification(config=snake_case__ )
							model.to(snake_case__ )
							model.eval()

							SCREAMING_SNAKE_CASE_   : List[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 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : List[str]					  =      self.num_choices
							SCREAMING_SNAKE_CASE_   : Any					  =      DistilBertForMultipleChoice(config=snake_case__ )
							model.to(snake_case__ )
							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(
							    snake_case__      ,attention_mask=snake_case__      ,labels=snake_case__      ,)
							self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.num_choices) )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      self.prepare_config_and_inputs()
							(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 				lowerCAmelCase_       (					__A						,     __A						,     unittest.TestCase							):
				__a     :       Optional[int]						=     (
				    (
				        DistilBertModel,
				        DistilBertForMaskedLM,
				        DistilBertForMultipleChoice,
				        DistilBertForQuestionAnswering,
				        DistilBertForSequenceClassification,
				        DistilBertForTokenClassification,
				    )
				    if is_torch_available()
				    else None
				)
				__a     :       Optional[Any]						=     (
				    {
				        "feature-extraction": DistilBertModel,
				        "fill-mask": DistilBertForMaskedLM,
				        "question-answering": DistilBertForQuestionAnswering,
				        "text-classification": DistilBertForSequenceClassification,
				        "token-classification": DistilBertForTokenClassification,
				        "zero-shot": DistilBertForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				__a     :       Optional[int]						=     True
				__a     :       List[str]						=     True
				__a     :       Union[str, Any]						=     True
				__a     :       str						=     True






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Dict					  =      DistilBertModelTester(self )
							SCREAMING_SNAKE_CASE_   : List[str]					  =      ConfigTester(self      ,config_class=snake_case__      ,dim=37 )






				def 		snake_case	(			self ):
							self.config_tester.run_common_tests()






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : int					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_model(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Dict					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_masked_lm(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_question_answering(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_sequence_classification(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : List[str]					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_token_classification(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_multiple_choice(*snake_case__ )






				@slow
				def 		snake_case	(			self ):

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






				@slow
				@require_torch_gpu
				def 		snake_case	(			self ):
							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=snake_case__ )

										SCREAMING_SNAKE_CASE_   : List[Any]					  =      self._prepare_for_class(snake_case__      ,snake_case__ )
										SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      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' ) )
													SCREAMING_SNAKE_CASE_   : Tuple					  =      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 				lowerCAmelCase_       (					unittest.TestCase							):






				@slow
				def 		snake_case	(			self ):
							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(snake_case__      ,attention_mask=snake_case__ )[0]
							SCREAMING_SNAKE_CASE_   : List[str]					  =      torch.Size((1, 11, 768) )
							self.assertEqual(output.shape      ,snake_case__ )
							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]      ,snake_case__      ,atol=1E-4 ) )






import unittest

import numpy as np
import torch

from diffusers import KarrasVePipeline, KarrasVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()






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







			@property
			def 				__UpperCamelCase			(      self				:			List[str]						)				->							Union[str, Any]:

				"""simple docstring"""

				torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   str					  =							UNetaDModel(
				    block_out_channels=(32, 64)   ,   layers_per_block=2   ,   sample_size=32   ,   in_channels=3   ,   out_channels=3   ,   down_block_types=("DownBlock2D", "AttnDownBlock2D")   ,   up_block_types=("AttnUpBlock2D", "UpBlock2D")   ,   )
				return model






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

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							self.dummy_uncond_unet
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							KarrasVeScheduler()

				SCREAMING_SNAKE_CASE :   Any					  =							KarrasVePipeline(unet=a   ,   scheduler=a						)
				pipe.to(a						)
				pipe.set_progress_bar_config(disable=a						)

				SCREAMING_SNAKE_CASE :   Any					  =							torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							pipe(num_inference_steps=2   ,   generator=a   ,   output_type="numpy"						).images

				SCREAMING_SNAKE_CASE :   List[str]					  =							torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   List[Any]					  =							pipe(num_inference_steps=2   ,   generator=a   ,   output_type="numpy"   ,   return_dict=a						)[0]

				SCREAMING_SNAKE_CASE :   List[Any]					  =							image[0, -3:, -3:, -1]
				SCREAMING_SNAKE_CASE :   Any					  =							image_from_tuple[0, -3:, -3:, -1]

				assert image.shape == (1, 32, 32, 3)
				SCREAMING_SNAKE_CASE :   str					  =							np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]						)

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






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







			def 				__UpperCamelCase			(      self				:			int						)				->							Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							"google/ncsnpp-celebahq-256"
				SCREAMING_SNAKE_CASE :   List[Any]					  =							UNetaDModel.from_pretrained(a						)
				SCREAMING_SNAKE_CASE :   Any					  =							KarrasVeScheduler()

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							KarrasVePipeline(unet=a   ,   scheduler=a						)
				pipe.to(a						)
				pipe.set_progress_bar_config(disable=a						)

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   Any					  =							pipe(num_inference_steps=20   ,   generator=a   ,   output_type="numpy"						).images

				SCREAMING_SNAKE_CASE :   List[str]					  =							image[0, -3:, -3:, -1]
				assert image.shape == (1, 256, 256, 3)
				SCREAMING_SNAKE_CASE :   str					  =							np.array([0.578, 0.5811, 0.5924, 0.5809, 0.587, 0.5886, 0.5861, 0.5802, 0.586]						)

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