





from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import tensorflow as tf

     from transformers.models.distilbert.modeling_tf_distilbert import (
         TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFDistilBertForMaskedLM,
         TFDistilBertForMultipleChoice,
         TFDistilBertForQuestionAnswering,
         TFDistilBertForSequenceClassification,
         TFDistilBertForTokenClassification,
         TFDistilBertModel,
     )



class   _SCREAMING_SNAKE_CASE		:







 '''simple docstring'''

 def __init__(self							:  Tuple			,					UpperCAmelCase_							:  List[Any]			,					)			->str:



   '''simple docstring'''



   lowerCamelCase__: Union[str, Any]         =parent
   lowerCamelCase__: Optional[Any]         =13
   lowerCamelCase__: List[str]         =7
   lowerCamelCase__: Any         =True
   lowerCamelCase__: List[Any]         =True
   lowerCamelCase__: Optional[int]         =False
   lowerCamelCase__: int         =True
   lowerCamelCase__: Optional[int]         =99
   lowerCamelCase__: Any         =32
   lowerCamelCase__: Tuple         =2
   lowerCamelCase__: Union[str, Any]         =4
   lowerCamelCase__: Union[str, Any]         =37
   lowerCamelCase__: List[Any]         ="gelu"
   lowerCamelCase__: Optional[int]         =0.1
   lowerCamelCase__: Optional[Any]         =0.1
   lowerCamelCase__: Optional[Any]         =512
   lowerCamelCase__: str         =16
   lowerCamelCase__: int         =2
   lowerCamelCase__: List[Any]         =0.02
   lowerCamelCase__: Optional[int]         =3
   lowerCamelCase__: Dict         =4
   lowerCamelCase__: Tuple         =None

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



   '''simple docstring'''



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

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

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

   lowerCamelCase__: str         =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			,					)

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

 def 				SCREAMING_SNAKE_CASE_   (self							:  str			,					UpperCAmelCase_							:  Dict			,					UpperCAmelCase_							:  str			,					UpperCAmelCase_							:  Union[str, Any]			,					UpperCAmelCase_							:  List[str]			,					UpperCAmelCase_							:  List[Any]			,					UpperCAmelCase_							:  Any)			->Dict:



   '''simple docstring'''



   lowerCamelCase__: List[Any]         =TFDistilBertModel(config=UpperCAmelCase_)
   lowerCamelCase__: Optional[int]         ={"input_ids": input_ids, "attention_mask": input_mask}

   lowerCamelCase__: List[str]         =model(UpperCAmelCase_)

   lowerCamelCase__: List[Any]         =[input_ids, input_mask]

   lowerCamelCase__: Any         =model(UpperCAmelCase_)

   self.parent.assertEqual(result.last_hidden_state.shape			,					(self.batch_size, self.seq_length, self.hidden_size))

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



   '''simple docstring'''



   lowerCamelCase__: List[str]         =TFDistilBertForMaskedLM(config=UpperCAmelCase_)
   lowerCamelCase__: Dict         ={"input_ids": input_ids, "attention_mask": input_mask}
   lowerCamelCase__: List[Any]         =model(UpperCAmelCase_)
   self.parent.assertEqual(result.logits.shape			,					(self.batch_size, self.seq_length, self.vocab_size))

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



   '''simple docstring'''



   lowerCamelCase__: List[str]         =TFDistilBertForQuestionAnswering(config=UpperCAmelCase_)
   lowerCamelCase__: Optional[Any]         ={
       "input_ids": input_ids,
       "attention_mask": input_mask,
   }
   lowerCamelCase__: int         =model(UpperCAmelCase_)
   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							:  List[Any]			,					UpperCAmelCase_							:  str			,					UpperCAmelCase_							:  int			,					UpperCAmelCase_							:  Dict			,					UpperCAmelCase_							:  Optional[int]			,					UpperCAmelCase_							:  Optional[int]			,					UpperCAmelCase_							:  Dict)			->int:



   '''simple docstring'''



   lowerCamelCase__: Optional[int]         =self.num_labels
   lowerCamelCase__: Optional[Any]         =TFDistilBertForSequenceClassification(UpperCAmelCase_)
   lowerCamelCase__: Optional[int]         ={"input_ids": input_ids, "attention_mask": input_mask}
   lowerCamelCase__: Union[str, Any]         =model(UpperCAmelCase_)
   self.parent.assertEqual(result.logits.shape			,					(self.batch_size, self.num_labels))

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



   '''simple docstring'''



   lowerCamelCase__: Dict         =self.num_choices
   lowerCamelCase__: Any         =TFDistilBertForMultipleChoice(UpperCAmelCase_)
   lowerCamelCase__: Optional[Any]         =tf.tile(tf.expand_dims(UpperCAmelCase_			,					1)			,					(1, self.num_choices, 1))
   lowerCamelCase__: Tuple         =tf.tile(tf.expand_dims(UpperCAmelCase_			,					1)			,					(1, self.num_choices, 1))
   lowerCamelCase__: Dict         ={
       "input_ids": multiple_choice_inputs_ids,
       "attention_mask": multiple_choice_input_mask,
   }
   lowerCamelCase__: Tuple         =model(UpperCAmelCase_)
   self.parent.assertEqual(result.logits.shape			,					(self.batch_size, self.num_choices))

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



   '''simple docstring'''



   lowerCamelCase__: Union[str, Any]         =self.num_labels
   lowerCamelCase__: List[str]         =TFDistilBertForTokenClassification(UpperCAmelCase_)
   lowerCamelCase__: str         ={"input_ids": input_ids, "attention_mask": input_mask}
   lowerCamelCase__: Optional[int]         =model(UpperCAmelCase_)
   self.parent.assertEqual(result.logits.shape			,					(self.batch_size, self.seq_length, self.num_labels))







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



   '''simple docstring'''



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



@require_tf
class   _SCREAMING_SNAKE_CASE		(      __SCREAMING_SNAKE_CASE       ,				__SCREAMING_SNAKE_CASE       ,				unittest.TestCase       ):







 '''simple docstring'''
 lowercase_						=     (
     (
         TFDistilBertModel,
         TFDistilBertForMaskedLM,
         TFDistilBertForQuestionAnswering,
         TFDistilBertForSequenceClassification,
         TFDistilBertForTokenClassification,
         TFDistilBertForMultipleChoice,
     )
     if is_tf_available()
     else None
 )
 lowercase_						=     (
     {
         "feature-extraction": TFDistilBertModel,
         "fill-mask": TFDistilBertForMaskedLM,
         "question-answering": TFDistilBertForQuestionAnswering,
         "text-classification": TFDistilBertForSequenceClassification,
         "token-classification": TFDistilBertForTokenClassification,
         "zero-shot": TFDistilBertForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 lowercase_						=     False
 lowercase_						=     False

 def 				SCREAMING_SNAKE_CASE_   (self							:  int)			->Any:



   '''simple docstring'''



   lowerCamelCase__: Tuple         =TFDistilBertModelTester(self)
   lowerCamelCase__: Dict         =ConfigTester(self			,					config_class=UpperCAmelCase_			,					dim=37)

 def 				SCREAMING_SNAKE_CASE_   (self							:  Optional[int])			->Tuple:



   '''simple docstring'''



   self.config_tester.run_common_tests()

 def 				SCREAMING_SNAKE_CASE_   (self							:  Any)			->int:



   '''simple docstring'''



   lowerCamelCase__: Any         =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_model(*UpperCAmelCase_)

 def 				SCREAMING_SNAKE_CASE_   (self							:  Tuple)			->Dict:



   '''simple docstring'''



   lowerCamelCase__: Union[str, Any]         =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_masked_lm(*UpperCAmelCase_)

 def 				SCREAMING_SNAKE_CASE_   (self							:  List[str])			->int:



   '''simple docstring'''



   lowerCamelCase__: Optional[Any]         =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_question_answering(*UpperCAmelCase_)

 def 				SCREAMING_SNAKE_CASE_   (self							:  List[Any])			->Tuple:



   '''simple docstring'''



   lowerCamelCase__: List[Any]         =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_sequence_classification(*UpperCAmelCase_)

 def 				SCREAMING_SNAKE_CASE_   (self							:  str)			->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__: Tuple         =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_multiple_choice(*UpperCAmelCase_)

 def 				SCREAMING_SNAKE_CASE_   (self							:  str)			->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__: int         =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_token_classification(*UpperCAmelCase_)







 @slow
 def 				SCREAMING_SNAKE_CASE_   (self							:  List[str])			->Dict:



   '''simple docstring'''



   for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]):
     lowerCamelCase__: str         =TFDistilBertModel.from_pretrained(UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)



@require_tf
class   _SCREAMING_SNAKE_CASE		(      unittest.TestCase       ):







 '''simple docstring'''

 @slow
 def 				SCREAMING_SNAKE_CASE_   (self							:  List[Any])			->Dict:



   '''simple docstring'''



   lowerCamelCase__: List[str]         =TFDistilBertModel.from_pretrained("distilbert-base-uncased")
   lowerCamelCase__: Any         =tf.constant([[0, 1, 2, 3, 4, 5]])
   lowerCamelCase__: Union[str, Any]         =model(UpperCAmelCase_)[0]

   lowerCamelCase__: Optional[int]         =[1, 6, 768]
   self.assertEqual(output.shape			,					UpperCAmelCase_)

   lowerCamelCase__: List[Any]         =tf.constant(
       [
           [
               [0.1926_1885, -0.1373_2955, 0.411_9799],
               [0.2215_0156, -0.0742_2661, 0.3903_7204],
               [0.2275_6018, -0.089_6414, 0.370_1467],
           ]
       ])
   tf.debugging.assert_near(output[:, :3, :3]			,					UpperCAmelCase_			,					atol=1E-4)








"""simple docstring"""






import unittest

import torch

from diffusers import DDIMScheduler, DDPMScheduler, UNetaDModel
from diffusers.training_utils import set_seed
from diffusers.utils.testing_utils import slow


__UpperCAmelCase						=			False



class 					__UpperCAmelCase							(						unittest.TestCase  ):
				def       UpperCAmelCase      (		self    :			Tuple	, a_    :			Optional[Any]=32    )	->	Union[str, Any]:



							'''simple docstring'''



							set_seed(0    )
							a__  :	Any		 =			UNetaDModel(sample_size=a_	, in_channels=3	, out_channels=3    )
							a__  :	Union[str, Any]		 =			torch.optim.SGD(model.parameters()	, lr=0.0001    )
							return model, optimizer






				@slow
				def       UpperCAmelCase      (		self    :			Union[str, Any]    )	->	Any:



							'''simple docstring'''



							a__  :	Optional[Any]		 =			"cpu"  # ensure full determinism without setting the CUBLAS_WORKSPACE_CONFIG env variable
							a__  :	Dict		 =			DDPMScheduler(
							    num_train_timesteps=10_00	, beta_start=0.0001	, beta_end=0.02	, beta_schedule="linear"	, clip_sample=a_	, )
							a__  :	Optional[Any]		 =			DDIMScheduler(
							    num_train_timesteps=10_00	, beta_start=0.0001	, beta_end=0.02	, beta_schedule="linear"	, clip_sample=a_	, )

							assert ddpm_scheduler.config.num_train_timesteps == ddim_scheduler.config.num_train_timesteps

							# shared batches for DDPM and DDIM
							set_seed(0    )
							a__  :	List[str]		 =			[torch.randn((4, 3, 32, 32)    ).clip(-1	, 1    ).to(a_    ) for _ in range(4    )]
							a__  :	List[str]		 =			[torch.randn((4, 3, 32, 32)    ).to(a_    ) for _ in range(4    )]
							a__  :	Optional[Any]		 =			[torch.randint(0	, 10_00	, (4,)    ).long().to(a_    ) for _ in range(4    )]

							# train with a DDPM scheduler
							a__ ,	a__  :	Union[str, Any]		 =			self.get_model_optimizer(resolution=32    )
							model.train().to(a_    )
							for i in range(4    ):
										optimizer.zero_grad()
										a__  :	Optional[int]		 =			ddpm_scheduler.add_noise(clean_images[i]	, noise[i]	, timesteps[i]    )
										a__  :	List[str]		 =			model(a_	, timesteps[i]    ).sample
										a__  :	Optional[int]		 =			torch.nn.functional.mse_loss(a_	, noise[i]    )
										loss.backward()
										optimizer.step()
							del model, optimizer

							# recreate the model and optimizer, and retry with DDIM
							a__ ,	a__  :	Optional[Any]		 =			self.get_model_optimizer(resolution=32    )
							model.train().to(a_    )
							for i in range(4    ):
										optimizer.zero_grad()
										a__  :	List[str]		 =			ddim_scheduler.add_noise(clean_images[i]	, noise[i]	, timesteps[i]    )
										a__  :	Optional[int]		 =			model(a_	, timesteps[i]    ).sample
										a__  :	Union[str, Any]		 =			torch.nn.functional.mse_loss(a_	, noise[i]    )
										loss.backward()
										optimizer.step()
							del model, optimizer

							self.assertTrue(torch.allclose(a_	, a_	, atol=1E-5    )    )
							self.assertTrue(torch.allclose(a_	, a_	, atol=1E-5    )    )