




'''simple docstring'''







from __future__ import annotations

import unittest

from transformers import CTRLConfig, 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.ctrl.modeling_tf_ctrl import (
						    TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFCTRLForSequenceClassification,
						    TFCTRLLMHeadModel,
						    TFCTRLModel,
						)





class      lowerCamelCase_							(       _a							):




					'''simple docstring'''






					def __init__(    self,  __UpperCamelCase,  ):
							_snake_case	       =       parent
							_snake_case	       =       13
							_snake_case	       =       7
							_snake_case	       =       True
							_snake_case	       =       True
							_snake_case	       =       True
							_snake_case	       =       True
							_snake_case	       =       True
							_snake_case	       =       99
							_snake_case	       =       32
							_snake_case	       =       2
							_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
							_snake_case	       =       self.vocab_size - 1






					def 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       ids_tensor([self.batch_size, self.seq_length],  self.vocab_size )

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

							_snake_case	       =       None
							if self.use_token_type_ids:
									_snake_case	       =       ids_tensor([self.batch_size, self.seq_length],  self.type_vocab_size )

							_snake_case	       =       None
							if self.use_mc_token_ids:
									_snake_case	       =       ids_tensor([self.batch_size, self.num_choices],  self.seq_length )

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

							_snake_case	       =       CTRLConfig(
							    vocab_size=self.vocab_size,  n_embd=self.hidden_size,  n_layer=self.num_hidden_layers,  n_head=self.num_attention_heads,  n_positions=self.max_position_embeddings,  pad_token_id=self.pad_token_id,  )

							_snake_case	       =       ids_tensor([self.num_hidden_layers, self.num_attention_heads],  2 )

							return (
							    config,
							    input_ids,
							    input_mask,
							    head_mask,
							    token_type_ids,
							    mc_token_ids,
							    sequence_labels,
							    token_labels,
							    choice_labels,
							)






					def 				__SCREAMING_SNAKE_CASE	(    self,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  *__UpperCamelCase ):
							_snake_case	       =       TFCTRLModel(config=_UpperCAmelCase )
							_snake_case	       =       {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
							_snake_case	       =       model(_UpperCAmelCase )

							_snake_case	       =       [input_ids, None, input_mask]  # None is the input for 'past'
							_snake_case	       =       model(_UpperCAmelCase )

							_snake_case	       =       model(_UpperCAmelCase )

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






					def 				__SCREAMING_SNAKE_CASE	(    self,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  *__UpperCamelCase ):
							_snake_case	       =       TFCTRLLMHeadModel(config=_UpperCAmelCase )
							_snake_case	       =       {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
							_snake_case	       =       model(_UpperCAmelCase )
							self.parent.assertEqual(result.logits.shape,  (self.batch_size, self.seq_length, self.vocab_size) )






					def 				__SCREAMING_SNAKE_CASE	(    self,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  *__UpperCamelCase ):
							_snake_case	       =       self.num_labels
							_snake_case	       =       ids_tensor([self.batch_size],  self.type_sequence_label_size )
							_snake_case	       =       {
							    """input_ids""": input_ids,
							    """token_type_ids""": token_type_ids,
							    """labels""": sequence_labels,
							}
							_snake_case	       =       TFCTRLForSequenceClassification(_UpperCAmelCase )
							_snake_case	       =       model(_UpperCAmelCase )
							self.parent.assertEqual(result.logits.shape,  (self.batch_size, self.num_labels) )






					def 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       self.prepare_config_and_inputs()

							(
							    _snake_case			
							)	       =       config_and_inputs

							_snake_case	       =       {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
							return config, inputs_dict





@require_tf
class      lowerCamelCase_							(       _a,       _a,       unittest.TestCase							):




					'''simple docstring'''






					lowercase__    = (TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
					lowercase__    = (TFCTRLLMHeadModel,) if is_tf_available() else ()
					lowercase__    = (
					    {
					        """feature-extraction""": TFCTRLModel,
					        """text-classification""": TFCTRLForSequenceClassification,
					        """text-generation""": TFCTRLLMHeadModel,
					        """zero-shot""": TFCTRLForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					lowercase__    = False
					lowercase__    = False






					def 				__SCREAMING_SNAKE_CASE	(    self,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase ):




							if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
									# Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
									# `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
									# config could not be created.
									return True

							return False






					def 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       TFCTRLModelTester(self )
							_snake_case	       =       ConfigTester(self,  config_class=_UpperCAmelCase,  n_embd=37 )






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






					def 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_ctrl_model(*_UpperCAmelCase )






					def 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_ctrl_lm_head(*_UpperCAmelCase )






					def 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_ctrl_for_sequence_classification(*_UpperCAmelCase )






					def 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       self.model_tester.prepare_config_and_inputs_for_common()
							_snake_case	       =       [TFCTRLLMHeadModel]
							_snake_case	       =       [TFCTRLForSequenceClassification]

							for model_class in self.all_model_classes:
									_snake_case	       =       model_class(_UpperCAmelCase )
									model.build()  # may be needed for the get_bias() call below
									assert isinstance(model.get_input_embeddings(),  tf.keras.layers.Layer )

									if model_class in list_lm_models:
											_snake_case	       =       model.get_output_embeddings()
											assert isinstance(_UpperCAmelCase,  tf.keras.layers.Layer )
											_snake_case	       =       model.get_bias()
											assert isinstance(_UpperCAmelCase,  _UpperCAmelCase )
											for k, v in name.items():
													assert isinstance(_UpperCAmelCase,  tf.Variable )
									elif model_class in list_other_models_with_output_ebd:
											_snake_case	       =       model.get_output_embeddings()
											assert isinstance(_UpperCAmelCase,  tf.keras.layers.Layer )
											_snake_case	       =       model.get_bias()
											assert name is None
									else:
											_snake_case	       =       model.get_output_embeddings()
											assert x is None
											_snake_case	       =       model.get_bias()
											assert name is None






					@slow
					def 				__SCREAMING_SNAKE_CASE	(    self ):




							for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									_snake_case	       =       TFCTRLModel.from_pretrained(_UpperCAmelCase )
									self.assertIsNotNone(_UpperCAmelCase )





@require_tf
class      lowerCamelCase_							(       unittest.TestCase							):




					'''simple docstring'''






					@slow
					def 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       TFCTRLLMHeadModel.from_pretrained("ctrl" )
							_snake_case	       =       tf.convert_to_tensor([[11_859, 0, 1_611, 8]],  dtype=tf.intaa )  # Legal the president is
							_snake_case	       =       [
							    11_859,
							    0,
							    1_611,
							    8,
							    5,
							    150,
							    26_449,
							    2,
							    19,
							    348,
							    469,
							    3,
							    2_595,
							    48,
							    20_740,
							    246_533,
							    246_533,
							    19,
							    30,
							    5,
							]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

							_snake_case	       =       model.generate(_UpperCAmelCase,  do_sample=_UpperCAmelCase )
							self.assertListEqual(output_ids[0].numpy().tolist(),  _UpperCAmelCase )











import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()
class 				__snake_case     (unittest.TestCase						):
				lowerCAmelCase__			=   "hf-internal-testing/tiny-stable-diffusion-pipe"
				lowerCAmelCase__			=   1
				lowerCAmelCase__			=   7_7
				lowerCAmelCase__			=   3_2
				lowerCAmelCase__			=   4
				lowerCAmelCase__			=   lowerCAmelCase__			=   6_4
				lowerCAmelCase__			=   torch.manual_seed(0						)



				def 					SCREAMING_SNAKE_CASE				(					self					:						List[str]						,		_UpperCAmelCase					:						Optional[Any]=0.9999					)     ->			List[Any]:





								'''simple docstring'''




								_lowerCAmelCase			:			List[Any]        =					UNetaDConditionModel.from_pretrained(self.model_id						,		subfolder="""unet"""					)
								_lowerCAmelCase			:			Optional[Any]        =					unet.to(_UpperCAmelCase					)
								_lowerCAmelCase			:			Optional[int]        =					EMAModel(unet.parameters()						,		decay=_UpperCAmelCase						,		model_cls=_UpperCAmelCase						,		model_config=unet.config					)
								return unet, ema_unet



				def 					SCREAMING_SNAKE_CASE				(					self					:						Any					)     ->			Optional[int]:





								'''simple docstring'''




								_lowerCAmelCase			:			Any        =					torch.randn(
								    self.batch_size						,		self.num_in_channels						,		self.latent_height						,		self.latent_width						,		generator=self.generator					).to(_UpperCAmelCase					)
								_lowerCAmelCase			:			List[str]        =					torch.randint(0						,		1000						,		size=(self.batch_size,)						,		generator=self.generator					).to(_UpperCAmelCase					)
								_lowerCAmelCase			:			Any        =					torch.randn(
								    self.batch_size						,		self.prompt_length						,		self.text_encoder_hidden_dim						,		generator=self.generator					).to(_UpperCAmelCase					)
								return noisy_latents, timesteps, encoder_hidden_states



				def 					SCREAMING_SNAKE_CASE				(					self					:						Dict						,		_UpperCAmelCase					:						Dict					)     ->			Dict:





								'''simple docstring'''




								_lowerCAmelCase			:			List[str]        =					{}
								for k, param in unet.state_dict().items():
												_lowerCAmelCase			:			List[Any]        =					torch.randn_like(_UpperCAmelCase					) + (param * torch.randn_like(_UpperCAmelCase					))
												updated_state_dict.update({k: updated_param}					)
								unet.load_state_dict(_UpperCAmelCase					)
								return unet



				def 					SCREAMING_SNAKE_CASE				(					self					:						Optional[Any]					)     ->			Tuple:





								'''simple docstring'''




								_lowerCAmelCase    ,			_lowerCAmelCase			:			Optional[int]        =					self.get_models()
								# Take the first (hypothetical) EMA step.
								ema_unet.step(unet.parameters()					)
								assert ema_unet.optimization_step == 1

								# Take two more.
								for _ in range(2					):
												ema_unet.step(unet.parameters()					)
								assert ema_unet.optimization_step == 3



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





								'''simple docstring'''




								_lowerCAmelCase    ,			_lowerCAmelCase			:			List[Any]        =					self.get_models()
								# Since the `unet` is not being updated (i.e., backprop'd)
								# there won't be any difference between the `params` of `unet`
								# and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
								ema_unet.step(unet.parameters()					)
								_lowerCAmelCase			:			Optional[int]        =					list(unet.parameters()					)
								for s_param, param in zip(ema_unet.shadow_params						,		_UpperCAmelCase					):
												assert torch.allclose(_UpperCAmelCase						,		_UpperCAmelCase					)

								# The above holds true even if we call `ema.step()` multiple times since
								# `unet` params are still not being updated.
								for _ in range(4					):
												ema_unet.step(unet.parameters()					)
								for s_param, param in zip(ema_unet.shadow_params						,		_UpperCAmelCase					):
												assert torch.allclose(_UpperCAmelCase						,		_UpperCAmelCase					)



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





								'''simple docstring'''




								_lowerCAmelCase    ,			_lowerCAmelCase			:			Union[str, Any]        =					self.get_models()
								# Here we simulate the parameter updates for `unet`. Since there might
								# be some parameters which are initialized to zero we take extra care to
								# initialize their values to something non-zero before the multiplication.
								_lowerCAmelCase			:			str        =					self.simulate_backprop(_UpperCAmelCase					)

								# Take the EMA step.
								ema_unet.step(unet_pseudo_updated_step_one.parameters()					)

								# Now the EMA'd parameters won't be equal to the original model parameters.
								_lowerCAmelCase			:			Tuple        =					list(unet_pseudo_updated_step_one.parameters()					)
								for s_param, param in zip(ema_unet.shadow_params						,		_UpperCAmelCase					):
												assert ~torch.allclose(_UpperCAmelCase						,		_UpperCAmelCase					)

								# Ensure this is the case when we take multiple EMA steps.
								for _ in range(4					):
												ema_unet.step(unet.parameters()					)
								for s_param, param in zip(ema_unet.shadow_params						,		_UpperCAmelCase					):
												assert ~torch.allclose(_UpperCAmelCase						,		_UpperCAmelCase					)



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





								'''simple docstring'''




								_lowerCAmelCase    ,			_lowerCAmelCase			:			str        =					self.get_models()

								# First backprop + EMA
								_lowerCAmelCase			:			Optional[int]        =					self.simulate_backprop(_UpperCAmelCase					)
								ema_unet.step(unet_step_one.parameters()					)
								_lowerCAmelCase			:			List[str]        =					ema_unet.shadow_params

								# Second backprop + EMA
								_lowerCAmelCase			:			Optional[int]        =					self.simulate_backprop(_UpperCAmelCase					)
								ema_unet.step(unet_step_two.parameters()					)
								_lowerCAmelCase			:			Tuple        =					ema_unet.shadow_params

								for step_one, step_two in zip(_UpperCAmelCase						,		_UpperCAmelCase					):
												assert ~torch.allclose(_UpperCAmelCase						,		_UpperCAmelCase					)



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





								'''simple docstring'''




								_lowerCAmelCase    ,			_lowerCAmelCase			:			int        =					self.get_models(decay=0.0					)
								_lowerCAmelCase			:			Dict        =					self.simulate_backprop(_UpperCAmelCase					)
								ema_unet.step(unet_step_one.parameters()					)
								_lowerCAmelCase			:			Any        =					ema_unet.shadow_params

								_lowerCAmelCase			:			List[Any]        =					self.simulate_backprop(_UpperCAmelCase					)
								ema_unet.step(unet_step_two.parameters()					)
								_lowerCAmelCase			:			int        =					ema_unet.shadow_params

								for step_one, step_two in zip(_UpperCAmelCase						,		_UpperCAmelCase					):
												assert torch.allclose(_UpperCAmelCase						,		_UpperCAmelCase					)



				@skip_mps
				def 					SCREAMING_SNAKE_CASE				(					self					:						Union[str, Any]					)     ->			str:





								'''simple docstring'''




								_lowerCAmelCase    ,			_lowerCAmelCase			:			Tuple        =					self.get_models()
								_lowerCAmelCase    ,			_lowerCAmelCase    ,			_lowerCAmelCase			:			Tuple        =					self.get_dummy_inputs()

								with tempfile.TemporaryDirectory() as tmpdir:
												ema_unet.save_pretrained(_UpperCAmelCase					)
												_lowerCAmelCase			:			Union[str, Any]        =					UNetaDConditionModel.from_pretrained(_UpperCAmelCase						,		model_cls=_UpperCAmelCase					)
												_lowerCAmelCase			:			Any        =					loaded_unet.to(unet.device					)

								# Since no EMA step has been performed the outputs should match.
								_lowerCAmelCase			:			str        =					unet(_UpperCAmelCase						,		_UpperCAmelCase						,		_UpperCAmelCase					).sample
								_lowerCAmelCase			:			Union[str, Any]        =					loaded_unet(_UpperCAmelCase						,		_UpperCAmelCase						,		_UpperCAmelCase					).sample

								assert torch.allclose(_UpperCAmelCase						,		_UpperCAmelCase						,		atol=1E-4					)

