

'''simple docstring'''
from __future__ import annotations

import unittest

from transformers import AlbertConfig, is_tf_available
from transformers.models.auto import get_values
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 import TF_MODEL_FOR_PRETRAINING_MAPPING
					from transformers.models.albert.modeling_tf_albert import (
					    TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
					    TFAlbertForMaskedLM,
					    TFAlbertForMultipleChoice,
					    TFAlbertForPreTraining,
					    TFAlbertForQuestionAnswering,
					    TFAlbertForSequenceClassification,
					    TFAlbertForTokenClassification,
					    TFAlbertModel,
					)
class  _snake_case						:






			def __init__(   self       ,_snake_case       ,_snake_case=13       ,_snake_case=7       ,_snake_case=True       ,_snake_case=True       ,_snake_case=True       ,_snake_case=True       ,_snake_case=99       ,_snake_case=16       ,_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=5_12       ,_snake_case=16       ,_snake_case=2       ,_snake_case=0.02       ,_snake_case=3       ,_snake_case=4       ,_snake_case=None       ,):
				UpperCAmelCase_			:	List[str]      =   parent
				UpperCAmelCase_			:	str      =   13
				UpperCAmelCase_			:	Optional[int]      =   7
				UpperCAmelCase_			:	int      =   True
				UpperCAmelCase_			:	str      =   True
				UpperCAmelCase_			:	Optional[Any]      =   True
				UpperCAmelCase_			:	List[str]      =   True
				UpperCAmelCase_			:	int      =   99
				UpperCAmelCase_			:	List[Any]      =   16
				UpperCAmelCase_			:	Dict      =   32
				UpperCAmelCase_			:	Dict      =   2
				UpperCAmelCase_			:	Dict      =   4
				UpperCAmelCase_			:	Optional[Any]      =   37
				UpperCAmelCase_			:	List[str]      =   "gelu"
				UpperCAmelCase_			:	Tuple      =   0.1
				UpperCAmelCase_			:	str      =   0.1
				UpperCAmelCase_			:	List[Any]      =   5_12
				UpperCAmelCase_			:	Tuple      =   16
				UpperCAmelCase_			:	Dict      =   2
				UpperCAmelCase_			:	List[Any]      =   0.02
				UpperCAmelCase_			:	Tuple      =   3
				UpperCAmelCase_			:	List[str]      =   4
				UpperCAmelCase_			:	Tuple      =   None






			def 						UpperCamelCase__				(   self			):
				UpperCAmelCase_			:	Union[str, Any]      =   ids_tensor([self.batch_size, self.seq_length]       ,self.vocab_size			)

				UpperCAmelCase_			:	Tuple      =   None
				if self.use_input_mask:
					UpperCAmelCase_			:	Optional[Any]      =   random_attention_mask([self.batch_size, self.seq_length]			)

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

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

				UpperCAmelCase_			:	Any      =   AlbertConfig(
				    vocab_size=self.vocab_size       ,hidden_size=self.hidden_size       ,embedding_size=self.embedding_size       ,num_hidden_layers=self.num_hidden_layers       ,num_attention_heads=self.num_attention_heads       ,intermediate_size=self.intermediate_size       ,hidden_act=self.hidden_act       ,hidden_dropout_prob=self.hidden_dropout_prob       ,attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,max_position_embeddings=self.max_position_embeddings       ,type_vocab_size=self.type_vocab_size       ,initializer_range=self.initializer_range       ,)

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






			def 						UpperCamelCase__				(   self       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case			):
				UpperCAmelCase_			:	Tuple      =   TFAlbertModel(config=_snake_case			)
				# inputs = {'input_ids': input_ids,
				#           'attention_mask': input_mask,
				#           'token_type_ids': token_type_ids}
				# sequence_output, pooled_output = model(**inputs)
				UpperCAmelCase_			:	str      =   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
				UpperCAmelCase_			:	List[str]      =   model(_snake_case			)

				UpperCAmelCase_			:	Dict      =   [input_ids, input_mask]
				UpperCAmelCase_			:	Dict      =   model(_snake_case			)

				UpperCAmelCase_			:	List[str]      =   model(_snake_case			)

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






			def 						UpperCamelCase__				(   self       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case			):
				UpperCAmelCase_			:	Any      =   self.num_labels
				UpperCAmelCase_			:	Any      =   TFAlbertForPreTraining(config=_snake_case			)
				UpperCAmelCase_			:	List[Any]      =   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
				UpperCAmelCase_			:	List[str]      =   model(_snake_case			)
				self.parent.assertEqual(result.prediction_logits.shape       ,(self.batch_size, self.seq_length, self.vocab_size)			)
				self.parent.assertEqual(result.sop_logits.shape       ,(self.batch_size, self.num_labels)			)






			def 						UpperCamelCase__				(   self       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case			):
				UpperCAmelCase_			:	Dict      =   TFAlbertForMaskedLM(config=_snake_case			)
				UpperCAmelCase_			:	Tuple      =   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
				UpperCAmelCase_			:	Dict      =   model(_snake_case			)
				self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.seq_length, self.vocab_size)			)






			def 						UpperCamelCase__				(   self       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case			):
				UpperCAmelCase_			:	List[Any]      =   self.num_labels
				UpperCAmelCase_			:	Union[str, Any]      =   TFAlbertForSequenceClassification(config=_snake_case			)
				UpperCAmelCase_			:	str      =   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
				UpperCAmelCase_			:	List[Any]      =   model(_snake_case			)
				self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.num_labels)			)






			def 						UpperCamelCase__				(   self       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case			):
				UpperCAmelCase_			:	Union[str, Any]      =   TFAlbertForQuestionAnswering(config=_snake_case			)
				UpperCAmelCase_			:	Union[str, Any]      =   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
				UpperCAmelCase_			:	Any      =   model(_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 						UpperCamelCase__				(   self       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case			):
				UpperCAmelCase_			:	Tuple      =   self.num_choices
				UpperCAmelCase_			:	List[str]      =   TFAlbertForMultipleChoice(config=_snake_case			)
				UpperCAmelCase_			:	List[Any]      =   tf.tile(tf.expand_dims(_snake_case       ,1			)       ,(1, self.num_choices, 1)			)
				UpperCAmelCase_			:	Union[str, Any]      =   tf.tile(tf.expand_dims(_snake_case       ,1			)       ,(1, self.num_choices, 1)			)
				UpperCAmelCase_			:	Optional[int]      =   tf.tile(tf.expand_dims(_snake_case       ,1			)       ,(1, self.num_choices, 1)			)
				UpperCAmelCase_			:	Dict      =   {
				    "input_ids": multiple_choice_inputs_ids,
				    "attention_mask": multiple_choice_input_mask,
				    "token_type_ids": multiple_choice_token_type_ids,
				}
				UpperCAmelCase_			:	Union[str, Any]      =   model(_snake_case			)
				self.parent.assertListEqual(list(result["logits"].shape			)       ,[self.batch_size, self.num_choices]			)






			def 						UpperCamelCase__				(   self       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case       ,_snake_case			):
				UpperCAmelCase_			:	List[Any]      =   self.num_labels
				UpperCAmelCase_			:	List[str]      =   TFAlbertForTokenClassification(config=_snake_case			)
				UpperCAmelCase_			:	int      =   {
				    "input_ids": input_ids,
				    "attention_mask": input_mask,
				    "token_type_ids": token_type_ids,
				}
				UpperCAmelCase_			:	List[str]      =   model(_snake_case			)
				self.parent.assertListEqual(list(result["logits"].shape			)       ,[self.batch_size, self.seq_length, self.num_labels]			)






			def 						UpperCamelCase__				(   self			):
				UpperCAmelCase_			:	List[Any]      =   self.prepare_config_and_inputs()
				(
				    (
				    UpperCAmelCase_
				) , (
				    UpperCAmelCase_
				) , (
				    UpperCAmelCase_
				) , (
				    UpperCAmelCase_
				) , (
				    UpperCAmelCase_
				) , (
				    UpperCAmelCase_
				) , (
				    UpperCAmelCase_
				) , 
				)			:	Union[str, Any]      =   config_and_inputs
				UpperCAmelCase_			:	Optional[int]      =   {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
				return config, inputs_dict




@require_tf
class  _snake_case						(__SCREAMING_SNAKE_CASE       ,       __SCREAMING_SNAKE_CASE       ,       unittest.TestCase):
			__A   :     Tuple          =(
			    (
			        TFAlbertModel,
			        TFAlbertForPreTraining,
			        TFAlbertForMaskedLM,
			        TFAlbertForSequenceClassification,
			        TFAlbertForQuestionAnswering,
			        TFAlbertForTokenClassification,
			        TFAlbertForMultipleChoice,
			    )
			    if is_tf_available()
			    else ()
			)
			__A   :     List[str]          =(
			    {
			        "feature-extraction": TFAlbertModel,
			        "fill-mask": TFAlbertForMaskedLM,
			        "question-answering": TFAlbertForQuestionAnswering,
			        "text-classification": TFAlbertForSequenceClassification,
			        "token-classification": TFAlbertForTokenClassification,
			        "zero-shot": TFAlbertForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			__A   :     List[str]          =False
			__A   :     Dict          =False






			def 						UpperCamelCase__				(   self       ,_snake_case       ,_snake_case       ,_snake_case=False			):
				UpperCAmelCase_			:	List[Any]      =   super()._prepare_for_class(_snake_case       ,_snake_case       ,return_labels=_snake_case			)

				if return_labels:
					if model_class in get_values(_snake_case			):
						UpperCAmelCase_			:	List[str]      =   tf.zeros(self.model_tester.batch_size       ,dtype=tf.intaa			)

				return inputs_dict






			def 						UpperCamelCase__				(   self			):
				UpperCAmelCase_			:	List[str]      =   TFAlbertModelTester(self			)
				UpperCAmelCase_			:	Dict      =   ConfigTester(self       ,config_class=_snake_case       ,hidden_size=37			)






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






			def 						UpperCamelCase__				(   self			):
				UpperCAmelCase_			:	Optional[Any]      =   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_albert_model(*_snake_case			)






			def 						UpperCamelCase__				(   self			):
				UpperCAmelCase_			:	Any      =   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_albert_for_pretraining(*_snake_case			)






			def 						UpperCamelCase__				(   self			):
				UpperCAmelCase_			:	Any      =   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_albert_for_masked_lm(*_snake_case			)






			def 						UpperCamelCase__				(   self			):
				UpperCAmelCase_			:	int      =   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_albert_for_multiple_choice(*_snake_case			)






			def 						UpperCamelCase__				(   self			):
				UpperCAmelCase_			:	Union[str, Any]      =   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_albert_for_sequence_classification(*_snake_case			)






			def 						UpperCamelCase__				(   self			):
				UpperCAmelCase_			:	Optional[int]      =   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_albert_for_question_answering(*_snake_case			)






			@slow
			def 						UpperCamelCase__				(   self			):
				for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					UpperCAmelCase_			:	Tuple      =   TFAlbertModel.from_pretrained(_snake_case			)
					self.assertIsNotNone(_snake_case			)




@require_tf
class  _snake_case						(unittest.TestCase):






			@slow
			def 						UpperCamelCase__				(   self			):
				UpperCAmelCase_			:	Optional[int]      =   TFAlbertForPreTraining.from_pretrained("albert-base-v2"			)
				UpperCAmelCase_			:	Tuple      =   tf.constant([[0, 1, 2, 3, 4, 5]]			)
				UpperCAmelCase_			:	int      =   model(_snake_case			)[0]

				UpperCAmelCase_			:	Union[str, Any]      =   [1, 6, 3_00_00]
				self.assertEqual(output.shape       ,_snake_case			)

				UpperCAmelCase_			:	List[Any]      =   tf.constant(
				    [
				        [
				            [4.595668, 0.74462754, -1.818147],
				            [4.5954347, 0.7454184, -1.8188258],
				            [4.5954905, 0.7448235, -1.8182316],
				        ]
				    ]			)
				tf.debugging.assert_near(output[:, :3, :3]       ,_snake_case       ,atol=1E-4			)







'''simple docstring'''
import torch

from diffusers import KDPMaAncestralDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest
class  _snake_case						(__SCREAMING_SNAKE_CASE):
			__A   :     Optional[int]          =(KDPMaAncestralDiscreteScheduler,)
			__A   :     List[Any]          =10






			def 						UpperCamelCase__				(   self       ,**_snake_case			):
				UpperCAmelCase_			:	List[Any]      =   {
				    "num_train_timesteps": 11_00,
				    "beta_start": 0.0001,
				    "beta_end": 0.02,
				    "beta_schedule": "linear",
				}

				config.update(**_snake_case			)
				return config






			def 						UpperCamelCase__				(   self			):
				for timesteps in [10, 50, 1_00, 10_00]:
					self.check_over_configs(num_train_timesteps=_snake_case			)






			def 						UpperCamelCase__				(   self			):
				for beta_start, beta_end in zip([0.00001, 0.0001, 0.001]       ,[0.0002, 0.002, 0.02]			):
					self.check_over_configs(beta_start=_snake_case       ,beta_end=_snake_case			)






			def 						UpperCamelCase__				(   self			):
				for schedule in ["linear", "scaled_linear"]:
					self.check_over_configs(beta_schedule=_snake_case			)






			def 						UpperCamelCase__				(   self			):
				if torch_device == "mps":
					return
				UpperCAmelCase_			:	str      =   self.scheduler_classes[0]
				UpperCAmelCase_			:	str      =   self.get_scheduler_config()
				UpperCAmelCase_			:	str      =   scheduler_class(**_snake_case			)

				scheduler.set_timesteps(self.num_inference_steps			)

				UpperCAmelCase_			:	Dict      =   torch.manual_seed(0			)

				UpperCAmelCase_			:	str      =   self.dummy_model()
				UpperCAmelCase_			:	str      =   self.dummy_sample_deter * scheduler.init_noise_sigma
				UpperCAmelCase_			:	Optional[Any]      =   sample.to(_snake_case			)

				for i, t in enumerate(scheduler.timesteps			):
					UpperCAmelCase_			:	Tuple      =   scheduler.scale_model_input(_snake_case       ,_snake_case			)

					UpperCAmelCase_			:	Union[str, Any]      =   model(_snake_case       ,_snake_case			)

					UpperCAmelCase_			:	Optional[Any]      =   scheduler.step(_snake_case       ,_snake_case       ,_snake_case       ,generator=_snake_case			)
					UpperCAmelCase_			:	Optional[Any]      =   output.prev_sample

				UpperCAmelCase_			:	Dict      =   torch.sum(torch.abs(_snake_case			)			)
				UpperCAmelCase_			:	List[Any]      =   torch.mean(torch.abs(_snake_case			)			)

				assert abs(result_sum.item() - 13849.3877			) < 1E-2
				assert abs(result_mean.item() - 18.0331			) < 5E-3






			def 						UpperCamelCase__				(   self			):
				for prediction_type in ["epsilon", "v_prediction"]:
					self.check_over_configs(prediction_type=_snake_case			)






			def 						UpperCamelCase__				(   self			):
				if torch_device == "mps":
					return
				UpperCAmelCase_			:	Any      =   self.scheduler_classes[0]
				UpperCAmelCase_			:	Any      =   self.get_scheduler_config(prediction_type="v_prediction"			)
				UpperCAmelCase_			:	Dict      =   scheduler_class(**_snake_case			)

				scheduler.set_timesteps(self.num_inference_steps			)

				UpperCAmelCase_			:	Union[str, Any]      =   self.dummy_model()
				UpperCAmelCase_			:	List[str]      =   self.dummy_sample_deter * scheduler.init_noise_sigma
				UpperCAmelCase_			:	int      =   sample.to(_snake_case			)

				UpperCAmelCase_			:	str      =   torch.manual_seed(0			)

				for i, t in enumerate(scheduler.timesteps			):
					UpperCAmelCase_			:	Dict      =   scheduler.scale_model_input(_snake_case       ,_snake_case			)

					UpperCAmelCase_			:	str      =   model(_snake_case       ,_snake_case			)

					UpperCAmelCase_			:	Optional[Any]      =   scheduler.step(_snake_case       ,_snake_case       ,_snake_case       ,generator=_snake_case			)
					UpperCAmelCase_			:	Any      =   output.prev_sample

				UpperCAmelCase_			:	Any      =   torch.sum(torch.abs(_snake_case			)			)
				UpperCAmelCase_			:	Tuple      =   torch.mean(torch.abs(_snake_case			)			)

				assert abs(result_sum.item() - 328.9970			) < 1E-2
				assert abs(result_mean.item() - 0.4284			) < 1E-3






			def 						UpperCamelCase__				(   self			):
				if torch_device == "mps":
					return
				UpperCAmelCase_			:	Optional[Any]      =   self.scheduler_classes[0]
				UpperCAmelCase_			:	Union[str, Any]      =   self.get_scheduler_config()
				UpperCAmelCase_			:	str      =   scheduler_class(**_snake_case			)

				scheduler.set_timesteps(self.num_inference_steps       ,device=_snake_case			)
				UpperCAmelCase_			:	Dict      =   torch.manual_seed(0			)

				UpperCAmelCase_			:	List[str]      =   self.dummy_model()
				UpperCAmelCase_			:	str      =   self.dummy_sample_deter.to(_snake_case			) * scheduler.init_noise_sigma

				for t in scheduler.timesteps:
					UpperCAmelCase_			:	Union[str, Any]      =   scheduler.scale_model_input(_snake_case       ,_snake_case			)

					UpperCAmelCase_			:	Tuple      =   model(_snake_case       ,_snake_case			)

					UpperCAmelCase_			:	int      =   scheduler.step(_snake_case       ,_snake_case       ,_snake_case       ,generator=_snake_case			)
					UpperCAmelCase_			:	List[str]      =   output.prev_sample

				UpperCAmelCase_			:	Any      =   torch.sum(torch.abs(_snake_case			)			)
				UpperCAmelCase_			:	Optional[int]      =   torch.mean(torch.abs(_snake_case			)			)

				assert abs(result_sum.item() - 13849.3818			) < 1E-1
				assert abs(result_mean.item() - 18.0331			) < 1E-3




