




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     __SCREAMING_SNAKE_CASE							(		lowerCamelCase  ):

					def __init__(    self			,	__lowercase			,	):







						'''simple docstring'''




						lowerCamelCase_							:		int							       =							parent
						lowerCamelCase_							:		Optional[Any]							       =							13
						lowerCamelCase_							:		Dict							       =							7
						lowerCamelCase_							:		Any							       =							True
						lowerCamelCase_							:		Union[str, Any]							       =							True
						lowerCamelCase_							:		str							       =							True
						lowerCamelCase_							:		Union[str, Any]							       =							True
						lowerCamelCase_							:		List[str]							       =							True
						lowerCamelCase_							:		str							       =							99
						lowerCamelCase_							:		Optional[Any]							       =							32
						lowerCamelCase_							:		Tuple							       =							2
						lowerCamelCase_							:		Optional[int]							       =							4
						lowerCamelCase_							:		List[str]							       =							37
						lowerCamelCase_							:		str							       =							'gelu'
						lowerCamelCase_							:		Union[str, Any]							       =							0.1
						lowerCamelCase_							:		str							       =							0.1
						lowerCamelCase_							:		Optional[int]							       =							512
						lowerCamelCase_							:		Union[str, Any]							       =							16
						lowerCamelCase_							:		Tuple							       =							2
						lowerCamelCase_							:		List[Any]							       =							0.02
						lowerCamelCase_							:		Any							       =							3
						lowerCamelCase_							:		List[str]							       =							4
						lowerCamelCase_							:		int							       =							None
						lowerCamelCase_							:		Optional[Any]							       =							self.vocab_size - 1

					def      lowercase_							(    self ):







						'''simple docstring'''




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

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

						lowerCamelCase_							:		List[Any]							       =							None
						if self.use_token_type_ids:
							lowerCamelCase_							:		str							       =							ids_tensor([self.batch_size, self.seq_length]			,	self.type_vocab_size )

						lowerCamelCase_							:		Dict							       =							None
						if self.use_mc_token_ids:
							lowerCamelCase_							:		Optional[Any]							       =							ids_tensor([self.batch_size, self.num_choices]			,	self.seq_length )

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

						lowerCamelCase_							:		Tuple							       =							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			,	)

						lowerCamelCase_							:		Optional[int]							       =							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      lowercase_							(    self			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	*__lowercase ):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[int]							       =							TFCTRLModel(config=__lowercase )
						lowerCamelCase_							:		List[Any]							       =							{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
						lowerCamelCase_							:		Optional[Any]							       =							model(__lowercase )

						lowerCamelCase_							:		Any							       =							[input_ids, None, input_mask]  # None is the input for 'past'
						lowerCamelCase_							:		Any							       =							model(__lowercase )

						lowerCamelCase_							:		Tuple							       =							model(__lowercase )

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

					def      lowercase_							(    self			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	*__lowercase ):







						'''simple docstring'''




						lowerCamelCase_							:		int							       =							TFCTRLLMHeadModel(config=__lowercase )
						lowerCamelCase_							:		str							       =							{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
						lowerCamelCase_							:		Optional[int]							       =							model(__lowercase )
						self.parent.assertEqual(result.logits.shape			,	(self.batch_size, self.seq_length, self.vocab_size) )

					def      lowercase_							(    self			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	*__lowercase ):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[Any]							       =							self.num_labels
						lowerCamelCase_							:		Any							       =							ids_tensor([self.batch_size]			,	self.type_sequence_label_size )
						lowerCamelCase_							:		List[Any]							       =							{
						    'input_ids': input_ids,
						    'token_type_ids': token_type_ids,
						    'labels': sequence_labels,
						}
						lowerCamelCase_							:		str							       =							TFCTRLForSequenceClassification(__lowercase )
						lowerCamelCase_							:		Union[str, Any]							       =							model(__lowercase )
						self.parent.assertEqual(result.logits.shape			,	(self.batch_size, self.num_labels) )





					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[Any]							       =							self.prepare_config_and_inputs()

						(
						    (
						    lowerCamelCase_
						)       ,      (
						    lowerCamelCase_
						)       ,      (
						    lowerCamelCase_
						)       ,      (
						    lowerCamelCase_
						)       ,      (
						    lowerCamelCase_
						)       ,      (
						    lowerCamelCase_
						)       ,      (
						    lowerCamelCase_
						)       ,      (
						    lowerCamelCase_
						)       ,      (
						    lowerCamelCase_
						)       ,      
						)							:		Any							       =							config_and_inputs

						lowerCamelCase_							:		Optional[Any]							       =							{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
						return config, inputs_dict


@require_tf
class     __SCREAMING_SNAKE_CASE							(		lowerCamelCase       ,					lowerCamelCase       ,					unittest.TestCase  ):
					_lowerCamelCase		  =			(TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
					_lowerCamelCase		  =			(TFCTRLLMHeadModel,) if is_tf_available() else ()
					_lowerCamelCase		  =			(
					    {
					        """feature-extraction""": TFCTRLModel,
					        """text-classification""": TFCTRLForSequenceClassification,
					        """text-generation""": TFCTRLLMHeadModel,
					        """zero-shot""": TFCTRLForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					_lowerCamelCase		  =			False
					_lowerCamelCase		  =			False

					def      lowercase_							(    self			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase			,	__lowercase ):







						'''simple docstring'''




						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      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Union[str, Any]							       =							TFCTRLModelTester(self )
						lowerCamelCase_							:		List[str]							       =							ConfigTester(self			,	config_class=__lowercase			,	n_embd=37 )

					def      lowercase_							(    self ):







						'''simple docstring'''




						self.config_tester.run_common_tests()

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Any							       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_ctrl_model(*__lowercase )

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[Any]							       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_ctrl_lm_head(*__lowercase )

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Union[str, Any]							       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_ctrl_for_sequence_classification(*__lowercase )

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_       ,      lowerCamelCase_							:		int							       =							self.model_tester.prepare_config_and_inputs_for_common()
						lowerCamelCase_							:		Union[str, Any]							       =							[TFCTRLLMHeadModel]
						lowerCamelCase_							:		Optional[Any]							       =							[TFCTRLForSequenceClassification]

						for model_class in self.all_model_classes:
							lowerCamelCase_							:		List[Any]							       =							model_class(__lowercase )
							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:
								lowerCamelCase_							:		int							       =							model.get_output_embeddings()
								assert isinstance(__lowercase			,	tf.keras.layers.Layer )
								lowerCamelCase_							:		Optional[int]							       =							model.get_bias()
								assert isinstance(__lowercase			,	__lowercase )
								for k, v in name.items():
									assert isinstance(__lowercase			,	tf.Variable )
							elif model_class in list_other_models_with_output_ebd:
								lowerCamelCase_							:		int							       =							model.get_output_embeddings()
								assert isinstance(__lowercase			,	tf.keras.layers.Layer )
								lowerCamelCase_							:		Any							       =							model.get_bias()
								assert name is None
							else:
								lowerCamelCase_							:		int							       =							model.get_output_embeddings()
								assert x is None
								lowerCamelCase_							:		List[str]							       =							model.get_bias()
								assert name is None





					@slow
					def      lowercase_							(    self ):







						'''simple docstring'''




						for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							lowerCamelCase_							:		Optional[int]							       =							TFCTRLModel.from_pretrained(__lowercase )
							self.assertIsNotNone(__lowercase )


@require_tf
class     __SCREAMING_SNAKE_CASE							(		unittest.TestCase  ):

					@slow
					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Dict							       =							TFCTRLLMHeadModel.from_pretrained('ctrl' )
						lowerCamelCase_							:		Optional[int]							       =							tf.convert_to_tensor([[11_859, 0, 1_611, 8]]			,	dtype=tf.intaa )  # Legal the president is
						lowerCamelCase_							:		List[Any]							       =							[
						    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

						lowerCamelCase_							:		Union[str, Any]							       =							model.generate(__lowercase			,	do_sample=__lowercase )
						self.assertListEqual(output_ids[0].numpy().tolist()			,	__lowercase )






import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest


class     __SCREAMING_SNAKE_CASE							(		lowerCamelCase  ):
					_lowerCamelCase		  =			(HeunDiscreteScheduler,)
					_lowerCamelCase		  =			10

					def      lowercase_							(    self			,	**__lowercase ):







						'''simple docstring'''




						lowerCamelCase_							:		List[str]							       =							{
						    'num_train_timesteps': 1_100,
						    'beta_start': 0.00_01,
						    'beta_end': 0.02,
						    'beta_schedule': 'linear',
						}

						config.update(**__lowercase )
						return config

					def      lowercase_							(    self ):







						'''simple docstring'''




						for timesteps in [10, 50, 100, 1_000]:
							self.check_over_configs(num_train_timesteps=__lowercase )

					def      lowercase_							(    self ):







						'''simple docstring'''




						for beta_start, beta_end in zip([0.0_00_01, 0.00_01, 0.0_01]			,	[0.00_02, 0.0_02, 0.02] ):
							self.check_over_configs(beta_start=__lowercase			,	beta_end=__lowercase )

					def      lowercase_							(    self ):







						'''simple docstring'''




						for schedule in ["linear", "scaled_linear", "exp"]:
							self.check_over_configs(beta_schedule=__lowercase )

					def      lowercase_							(    self ):







						'''simple docstring'''




						for clip_sample_range in [1.0, 2.0, 3.0]:
							self.check_over_configs(clip_sample_range=__lowercase			,	clip_sample=__lowercase )

					def      lowercase_							(    self ):







						'''simple docstring'''




						for prediction_type in ["epsilon", "v_prediction", "sample"]:
							self.check_over_configs(prediction_type=__lowercase )

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Tuple							       =							self.scheduler_classes[0]
						lowerCamelCase_							:		List[str]							       =							self.get_scheduler_config()
						lowerCamelCase_							:		Union[str, Any]							       =							scheduler_class(**__lowercase )

						scheduler.set_timesteps(self.num_inference_steps )

						lowerCamelCase_							:		Optional[Any]							       =							self.dummy_model()
						lowerCamelCase_							:		Optional[int]							       =							self.dummy_sample_deter * scheduler.init_noise_sigma
						lowerCamelCase_							:		Dict							       =							sample.to(__lowercase )

						for i, t in enumerate(scheduler.timesteps ):
							lowerCamelCase_							:		Any							       =							scheduler.scale_model_input(__lowercase			,	__lowercase )

							lowerCamelCase_							:		int							       =							model(__lowercase			,	__lowercase )

							lowerCamelCase_							:		Optional[int]							       =							scheduler.step(__lowercase			,	__lowercase			,	__lowercase )
							lowerCamelCase_							:		Tuple							       =							output.prev_sample

						lowerCamelCase_							:		str							       =							torch.sum(torch.abs(__lowercase ) )
						lowerCamelCase_							:		Tuple							       =							torch.mean(torch.abs(__lowercase ) )

						if torch_device in ["cpu", "mps"]:
							assert abs(result_sum.item() - 0.12_33 ) < 1E-2
							assert abs(result_mean.item() - 0.00_02 ) < 1E-3
						else:
							# CUDA
							assert abs(result_sum.item() - 0.12_33 ) < 1E-2
							assert abs(result_mean.item() - 0.00_02 ) < 1E-3

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Any							       =							self.scheduler_classes[0]
						lowerCamelCase_							:		str							       =							self.get_scheduler_config(prediction_type='v_prediction' )
						lowerCamelCase_							:		Optional[Any]							       =							scheduler_class(**__lowercase )

						scheduler.set_timesteps(self.num_inference_steps )

						lowerCamelCase_							:		Any							       =							self.dummy_model()
						lowerCamelCase_							:		str							       =							self.dummy_sample_deter * scheduler.init_noise_sigma
						lowerCamelCase_							:		List[Any]							       =							sample.to(__lowercase )

						for i, t in enumerate(scheduler.timesteps ):
							lowerCamelCase_							:		int							       =							scheduler.scale_model_input(__lowercase			,	__lowercase )

							lowerCamelCase_							:		Tuple							       =							model(__lowercase			,	__lowercase )

							lowerCamelCase_							:		int							       =							scheduler.step(__lowercase			,	__lowercase			,	__lowercase )
							lowerCamelCase_							:		Dict							       =							output.prev_sample

						lowerCamelCase_							:		Any							       =							torch.sum(torch.abs(__lowercase ) )
						lowerCamelCase_							:		List[Any]							       =							torch.mean(torch.abs(__lowercase ) )

						if torch_device in ["cpu", "mps"]:
							assert abs(result_sum.item() - 4.6934E-07 ) < 1E-2
							assert abs(result_mean.item() - 6.1112E-10 ) < 1E-3
						else:
							# CUDA
							assert abs(result_sum.item() - 4.693428650170972E-07 ) < 1E-2
							assert abs(result_mean.item() - 0.00_02 ) < 1E-3

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Union[str, Any]							       =							self.scheduler_classes[0]
						lowerCamelCase_							:		int							       =							self.get_scheduler_config()
						lowerCamelCase_							:		Dict							       =							scheduler_class(**__lowercase )

						scheduler.set_timesteps(self.num_inference_steps			,	device=__lowercase )

						lowerCamelCase_							:		str							       =							self.dummy_model()
						lowerCamelCase_							:		Optional[int]							       =							self.dummy_sample_deter.to(__lowercase ) * scheduler.init_noise_sigma

						for t in scheduler.timesteps:
							lowerCamelCase_							:		Tuple							       =							scheduler.scale_model_input(__lowercase			,	__lowercase )

							lowerCamelCase_							:		Tuple							       =							model(__lowercase			,	__lowercase )

							lowerCamelCase_							:		List[Any]							       =							scheduler.step(__lowercase			,	__lowercase			,	__lowercase )
							lowerCamelCase_							:		Union[str, Any]							       =							output.prev_sample

						lowerCamelCase_							:		Dict							       =							torch.sum(torch.abs(__lowercase ) )
						lowerCamelCase_							:		Any							       =							torch.mean(torch.abs(__lowercase ) )

						if str(__lowercase ).startswith('cpu' ):
							# The following sum varies between 148 and 156 on mps. Why?
							assert abs(result_sum.item() - 0.12_33 ) < 1E-2
							assert abs(result_mean.item() - 0.00_02 ) < 1E-3
						elif str(__lowercase ).startswith('mps' ):
							# Larger tolerance on mps
							assert abs(result_mean.item() - 0.00_02 ) < 1E-2
						else:
							# CUDA
							assert abs(result_sum.item() - 0.12_33 ) < 1E-2
							assert abs(result_mean.item() - 0.00_02 ) < 1E-3





					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Union[str, Any]							       =							self.scheduler_classes[0]
						lowerCamelCase_							:		Optional[Any]							       =							self.get_scheduler_config()
						lowerCamelCase_							:		Dict							       =							scheduler_class(**__lowercase			,	use_karras_sigmas=__lowercase )

						scheduler.set_timesteps(self.num_inference_steps			,	device=__lowercase )

						lowerCamelCase_							:		int							       =							self.dummy_model()
						lowerCamelCase_							:		Optional[Any]							       =							self.dummy_sample_deter.to(__lowercase ) * scheduler.init_noise_sigma
						lowerCamelCase_							:		Tuple							       =							sample.to(__lowercase )

						for t in scheduler.timesteps:
							lowerCamelCase_							:		str							       =							scheduler.scale_model_input(__lowercase			,	__lowercase )

							lowerCamelCase_							:		Tuple							       =							model(__lowercase			,	__lowercase )

							lowerCamelCase_							:		Dict							       =							scheduler.step(__lowercase			,	__lowercase			,	__lowercase )
							lowerCamelCase_							:		Dict							       =							output.prev_sample

						lowerCamelCase_							:		Tuple							       =							torch.sum(torch.abs(__lowercase ) )
						lowerCamelCase_							:		Optional[Any]							       =							torch.mean(torch.abs(__lowercase ) )

						assert abs(result_sum.item() - 0.0_00_15 ) < 1E-2
						assert abs(result_mean.item() - 1.9869554535034695E-07 ) < 1E-2
