
'''simple docstring'''






from __future__ import annotations

import unittest

from transformers import 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 numpy
			import tensorflow as tf

			from transformers import (
			    TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
			    TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
			    TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
			    BertConfig,
			    DPRConfig,
			    TFDPRContextEncoder,
			    TFDPRQuestionEncoder,
			    TFDPRReader,
			)




class     __snake_case		:
					def __init__(     self,    A,    A=13,    A=7,    A=True,    A=True,    A=True,    A=True,    A=99,    A=32,    A=2,    A=4,    A=37,    A="gelu",    A=0.1,    A=0.1,    A=512,    A=16,    A=2,    A=0.02,    A=3,    A=4,    A=None,    A=0,    ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[Any]									= parent
									lowerCamelCase					:			Any									= batch_size
									lowerCamelCase					:			Optional[Any]									= seq_length
									lowerCamelCase					:			List[str]									= is_training
									lowerCamelCase					:			Optional[int]									= use_input_mask
									lowerCamelCase					:			int									= use_token_type_ids
									lowerCamelCase					:			List[Any]									= use_labels
									lowerCamelCase					:			List[str]									= vocab_size
									lowerCamelCase					:			List[Any]									= hidden_size
									lowerCamelCase					:			Dict									= num_hidden_layers
									lowerCamelCase					:			List[str]									= num_attention_heads
									lowerCamelCase					:			Union[str, Any]									= intermediate_size
									lowerCamelCase					:			int									= hidden_act
									lowerCamelCase					:			str									= hidden_dropout_prob
									lowerCamelCase					:			Union[str, Any]									= attention_probs_dropout_prob
									lowerCamelCase					:			str									= max_position_embeddings
									lowerCamelCase					:			int									= type_vocab_size
									lowerCamelCase					:			List[Any]									= type_sequence_label_size
									lowerCamelCase					:			int									= initializer_range
									lowerCamelCase					:			Any									= num_labels
									lowerCamelCase					:			Tuple									= num_choices
									lowerCamelCase					:			Dict									= scope
									lowerCamelCase					:			List[Any]									= projection_dim
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Tuple									= ids_tensor([self.batch_size, self.seq_length],    self.vocab_size      )

									lowerCamelCase					:			Tuple									= None
									if self.use_input_mask:
													# follow test_modeling_tf_ctrl.py
													lowerCamelCase					:			int									= random_attention_mask([self.batch_size, self.seq_length]      )

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

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

									lowerCamelCase					:			int									= BertConfig(
									    vocab_size=self.vocab_size,    hidden_size=self.hidden_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,    is_decoder=A,    initializer_range=self.initializer_range,    )
									lowerCamelCase					:			Dict									= DPRConfig(projection_dim=self.projection_dim,    **config.to_dict()      )

									return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
					def     UpperCAmelCase_							(     self,    A,    A,    A,    A,    A,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			Dict									= TFDPRContextEncoder(config=A      )
									lowerCamelCase					:			Optional[Any]									= model(A,    attention_mask=A,    token_type_ids=A      )
									lowerCamelCase					:			int									= model(A,    token_type_ids=A      )
									lowerCamelCase					:			int									= model(A      )
									self.parent.assertEqual(result.pooler_output.shape,    (self.batch_size, self.projection_dim or self.hidden_size)      )
					def     UpperCAmelCase_							(     self,    A,    A,    A,    A,    A,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= TFDPRQuestionEncoder(config=A      )
									lowerCamelCase					:			Optional[int]									= model(A,    attention_mask=A,    token_type_ids=A      )
									lowerCamelCase					:			Union[str, Any]									= model(A,    token_type_ids=A      )
									lowerCamelCase					:			List[Any]									= model(A      )
									self.parent.assertEqual(result.pooler_output.shape,    (self.batch_size, self.projection_dim or self.hidden_size)      )
					def     UpperCAmelCase_							(     self,    A,    A,    A,    A,    A,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			Dict									= TFDPRReader(config=A      )
									lowerCamelCase					:			Dict									= model(A,    attention_mask=A      )

									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)      )
									self.parent.assertEqual(result.relevance_logits.shape,    (self.batch_size,)      )




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			int									= self.prepare_config_and_inputs()
									(
									    lowerCamelCase 
									)					:			Dict									= config_and_inputs
									lowerCamelCase					:			Optional[int]									= {'input_ids': input_ids}
									return config, inputs_dict





@require_tf
class     __snake_case		(   a__    ,   a__    ,   unittest.TestCase):
					_lowerCAmelCase			     =	(
					    (
					        TFDPRContextEncoder,
					        TFDPRQuestionEncoder,
					        TFDPRReader,
					    )
					    if is_tf_available()
					    else ()
					)
					_lowerCAmelCase			     =	{'''feature-extraction''': TFDPRQuestionEncoder} if is_tf_available() else {}

					_lowerCAmelCase			     =	False
					_lowerCAmelCase			     =	False
					_lowerCAmelCase			     =	False
					_lowerCAmelCase			     =	False
					_lowerCAmelCase			     =	False
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Dict									= TFDPRModelTester(self      )
									lowerCamelCase					:			Union[str, Any]									= ConfigTester(self,    config_class=A,    hidden_size=37      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


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

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_dpr_context_encoder(*A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			int									= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_dpr_question_encoder(*A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Tuple									= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_dpr_reader(*A      )




					@slow
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													lowerCamelCase					:			Any									= TFDPRContextEncoder.from_pretrained(A      )
													self.assertIsNotNone(A      )

									for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													lowerCamelCase					:			Union[str, Any]									= TFDPRContextEncoder.from_pretrained(A      )
													self.assertIsNotNone(A      )

									for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													lowerCamelCase					:			str									= TFDPRQuestionEncoder.from_pretrained(A      )
													self.assertIsNotNone(A      )

									for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													lowerCamelCase					:			Union[str, Any]									= TFDPRReader.from_pretrained(A      )
													self.assertIsNotNone(A      )





@require_tf
class     __snake_case		(   unittest.TestCase):
					@slow
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Tuple									= TFDPRQuestionEncoder.from_pretrained('facebook/dpr-question_encoder-single-nq-base'      )

									lowerCamelCase					:			Optional[int]									= tf.constant(
									    [[101, 7592, 1010, 2003, 2026, 3899, 1_0140, 1029, 102]]      )  # [CLS] hello, is my dog cute? [SEP]
									lowerCamelCase					:			Optional[Any]									= model(A      )[0]  # embedding shape = (1, 768)
									# compare the actual values for a slice.
									lowerCamelCase					:			List[Any]									= tf.constant(
									    [
									        [
									            0.0323_6253,
									            0.1275_3335,
									            0.1681_8509,
									            0.0027_9786,
									            0.389_6933,
									            0.2426_4945,
									            0.217_8971,
									            -0.0233_5227,
									            -0.0848_1959,
									            -0.1432_4117,
									        ]
									    ]      )
									self.assertTrue(numpy.allclose(output[:, :10].numpy(),    expected_slice.numpy(),    atol=1e-4      )      )


'''simple docstring'''






import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest




class     __snake_case		(   a__):
					_lowerCAmelCase			     =	(DPMSolverSinglestepScheduler,)
					_lowerCAmelCase			     =	(('''num_inference_steps''', 25),)
					def     UpperCAmelCase_							(     self,    **A      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= {
									    'num_train_timesteps': 1000,
									    'beta_start': 0.0001,
									    'beta_end': 0.02,
									    'beta_schedule': 'linear',
									    'solver_order': 2,
									    'prediction_type': 'epsilon',
									    'thresholding': False,
									    'sample_max_value': 1.0,
									    'algorithm_type': 'dpmsolver++',
									    'solver_type': 'midpoint',
									    'lambda_min_clipped': -float('inf'      ),
									    'variance_type': None,
									}

									config.update(**A      )
									return config
					def     UpperCAmelCase_							(     self,    A=0,    **A      ):

									"""simple docstring"""


									lowerCamelCase					:			List[str]									= dict(self.forward_default_kwargs      )
									lowerCamelCase					:			Optional[Any]									= kwargs.pop('num_inference_steps',    A      )
									lowerCamelCase					:			Union[str, Any]									= self.dummy_sample
									lowerCamelCase					:			Dict									= 0.1 * sample
									lowerCamelCase					:			Dict									= [residual + 0.2, residual + 0.15, residual + 0.10]

									for scheduler_class in self.scheduler_classes:
													lowerCamelCase					:			Optional[Any]									= self.get_scheduler_config(**A      )
													lowerCamelCase					:			Dict									= scheduler_class(**A      )
													scheduler.set_timesteps(A      )
													# copy over dummy past residuals
													lowerCamelCase					:			str									= dummy_past_residuals[: scheduler.config.solver_order]

													with tempfile.TemporaryDirectory() as tmpdirname:
																	scheduler.save_config(A      )
																	lowerCamelCase					:			List[Any]									= scheduler_class.from_pretrained(A      )
																	new_scheduler.set_timesteps(A      )
																	# copy over dummy past residuals
																	lowerCamelCase					:			Optional[int]									= dummy_past_residuals[: new_scheduler.config.solver_order]

													lowerCamelCase			, lowerCamelCase					:			Optional[int]									= sample, sample
													for t in range(A,    time_step + scheduler.config.solver_order + 1      ):
																	lowerCamelCase					:			Dict									= scheduler.step(A,    A,    A,    **A      ).prev_sample
																	lowerCamelCase					:			Optional[int]									= new_scheduler.step(A,    A,    A,    **A      ).prev_sample

																	assert torch.sum(torch.abs(output - new_output      )      ) < 1e-5, "Scheduler outputs are not identical"
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									pass
					def     UpperCAmelCase_							(     self,    A=0,    **A      ):

									"""simple docstring"""


									lowerCamelCase					:			List[str]									= dict(self.forward_default_kwargs      )
									lowerCamelCase					:			str									= kwargs.pop('num_inference_steps',    A      )
									lowerCamelCase					:			Union[str, Any]									= self.dummy_sample
									lowerCamelCase					:			List[str]									= 0.1 * sample
									lowerCamelCase					:			List[str]									= [residual + 0.2, residual + 0.15, residual + 0.10]

									for scheduler_class in self.scheduler_classes:
													lowerCamelCase					:			Tuple									= self.get_scheduler_config()
													lowerCamelCase					:			Optional[Any]									= scheduler_class(**A      )
													scheduler.set_timesteps(A      )

													# copy over dummy past residuals (must be after setting timesteps)
													lowerCamelCase					:			Any									= dummy_past_residuals[: scheduler.config.solver_order]

													with tempfile.TemporaryDirectory() as tmpdirname:
																	scheduler.save_config(A      )
																	lowerCamelCase					:			Tuple									= scheduler_class.from_pretrained(A      )
																	# copy over dummy past residuals
																	new_scheduler.set_timesteps(A      )

																	# copy over dummy past residual (must be after setting timesteps)
																	lowerCamelCase					:			List[Any]									= dummy_past_residuals[: new_scheduler.config.solver_order]

													lowerCamelCase					:			int									= scheduler.step(A,    A,    A,    **A      ).prev_sample
													lowerCamelCase					:			Dict									= new_scheduler.step(A,    A,    A,    **A      ).prev_sample

													assert torch.sum(torch.abs(output - new_output      )      ) < 1e-5, "Scheduler outputs are not identical"
					def     UpperCAmelCase_							(     self,    A=None,    **A      ):

									"""simple docstring"""


									if scheduler is None:
													lowerCamelCase					:			Any									= self.scheduler_classes[0]
													lowerCamelCase					:			Optional[Any]									= self.get_scheduler_config(**A      )
													lowerCamelCase					:			Optional[int]									= scheduler_class(**A      )

									lowerCamelCase					:			List[Any]									= self.scheduler_classes[0]
									lowerCamelCase					:			Optional[Any]									= self.get_scheduler_config(**A      )
									lowerCamelCase					:			Optional[int]									= scheduler_class(**A      )

									lowerCamelCase					:			Any									= 10
									lowerCamelCase					:			Dict									= self.dummy_model()
									lowerCamelCase					:			Any									= self.dummy_sample_deter
									scheduler.set_timesteps(A      )

									for i, t in enumerate(scheduler.timesteps      ):
													lowerCamelCase					:			Dict									= model(A,    A      )
													lowerCamelCase					:			List[str]									= scheduler.step(A,    A,    A      ).prev_sample

									return sample
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Dict									= DPMSolverSinglestepScheduler(**self.get_scheduler_config()      )
									lowerCamelCase					:			Dict									= 50
									lowerCamelCase					:			Tuple									= self.dummy_model()
									lowerCamelCase					:			Optional[int]									= self.dummy_sample_deter
									scheduler.set_timesteps(A      )

									# make sure that the first t is uneven
									for i, t in enumerate(scheduler.timesteps[3:]      ):
													lowerCamelCase					:			Any									= model(A,    A      )
													lowerCamelCase					:			Optional[int]									= scheduler.step(A,    A,    A      ).prev_sample

									lowerCamelCase					:			Any									= torch.mean(torch.abs(A      )      )

									assert abs(result_mean.item() - 0.2574      ) < 1e-3
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									for timesteps in [25, 50, 100, 999, 1000]:
													self.check_over_configs(num_train_timesteps=A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Dict									= DPMSolverSinglestepScheduler(**self.get_scheduler_config()      )
									lowerCamelCase					:			str									= self.full_loop(scheduler=A      )
									lowerCamelCase					:			Optional[int]									= torch.mean(torch.abs(A      )      )

									assert abs(result_mean.item() - 0.2791      ) < 1e-3

									lowerCamelCase					:			Dict									= DEISMultistepScheduler.from_config(scheduler.config      )
									lowerCamelCase					:			Optional[int]									= DPMSolverMultistepScheduler.from_config(scheduler.config      )
									lowerCamelCase					:			Any									= UniPCMultistepScheduler.from_config(scheduler.config      )
									lowerCamelCase					:			Optional[Any]									= DPMSolverSinglestepScheduler.from_config(scheduler.config      )

									lowerCamelCase					:			str									= self.full_loop(scheduler=A      )
									lowerCamelCase					:			Optional[int]									= torch.mean(torch.abs(A      )      )

									assert abs(result_mean.item() - 0.2791      ) < 1e-3
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									self.check_over_configs(thresholding=A      )
									for order in [1, 2, 3]:
													for solver_type in ["midpoint", "heun"]:
																	for threshold in [0.5, 1.0, 2.0]:
																					for prediction_type in ["epsilon", "sample"]:
																									self.check_over_configs(
																									    thresholding=A,    prediction_type=A,    sample_max_value=A,    algorithm_type='dpmsolver++',    solver_order=A,    solver_type=A,    )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


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

									"""simple docstring"""


									for algorithm_type in ["dpmsolver", "dpmsolver++"]:
													for solver_type in ["midpoint", "heun"]:
																	for order in [1, 2, 3]:
																					for prediction_type in ["epsilon", "sample"]:
																									self.check_over_configs(
																									    solver_order=A,    solver_type=A,    prediction_type=A,    algorithm_type=A,    )
																									lowerCamelCase					:			Optional[Any]									= self.full_loop(
																									    solver_order=A,    solver_type=A,    prediction_type=A,    algorithm_type=A,    )
																									assert not torch.isnan(A      ).any(), "Samples have nan numbers"
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									self.check_over_configs(lower_order_final=A      )
									self.check_over_configs(lower_order_final=A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									self.check_over_configs(lambda_min_clipped=-float('inf'      )      )
									self.check_over_configs(lambda_min_clipped=-5.1      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									self.check_over_configs(variance_type=A      )
									self.check_over_configs(variance_type='learned_range'      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									for num_inference_steps in [1, 2, 3, 5, 10, 50, 100, 999, 1000]:
													self.check_over_forward(num_inference_steps=A,    time_step=0      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= self.full_loop()
									lowerCamelCase					:			str									= torch.mean(torch.abs(A      )      )

									assert abs(result_mean.item() - 0.2791      ) < 1e-3
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= self.full_loop(use_karras_sigmas=A      )
									lowerCamelCase					:			Tuple									= torch.mean(torch.abs(A      )      )

									assert abs(result_mean.item() - 0.2248      ) < 1e-3
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= self.full_loop(prediction_type='v_prediction'      )
									lowerCamelCase					:			Dict									= torch.mean(torch.abs(A      )      )

									assert abs(result_mean.item() - 0.1453      ) < 1e-3
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= self.full_loop(prediction_type='v_prediction',    use_karras_sigmas=A      )
									lowerCamelCase					:			Optional[Any]									= torch.mean(torch.abs(A      )      )

									assert abs(result_mean.item() - 0.0649      ) < 1e-3




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[Any]									= self.scheduler_classes[0]
									lowerCamelCase					:			Dict									= self.get_scheduler_config(thresholding=A,    dynamic_thresholding_ratio=0      )
									lowerCamelCase					:			str									= scheduler_class(**A      )

									lowerCamelCase					:			List[Any]									= 10
									lowerCamelCase					:			List[str]									= self.dummy_model()
									lowerCamelCase					:			int									= self.dummy_sample_deter.half()
									scheduler.set_timesteps(A      )

									for i, t in enumerate(scheduler.timesteps      ):
													lowerCamelCase					:			str									= model(A,    A      )
													lowerCamelCase					:			Tuple									= scheduler.step(A,    A,    A      ).prev_sample

									assert sample.dtype == torch.floataa
