



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       UpperCAmelCase__:


   def __init__(						self  :			Optional[Any]			,				lowerCAmelCase  :			Optional[Any]			,				lowerCAmelCase  :			List[Any]=13			,				lowerCAmelCase  :			List[str]=7			,				lowerCAmelCase  :			Optional[Any]=True			,				lowerCAmelCase  :			str=True			,				lowerCAmelCase  :			List[str]=True			,				lowerCAmelCase  :			str=True			,				lowerCAmelCase  :			List[Any]=99			,				lowerCAmelCase  :			Any=32			,				lowerCAmelCase  :			List[str]=2			,				lowerCAmelCase  :			int=4			,				lowerCAmelCase  :			str=37			,				lowerCAmelCase  :			Any="gelu"			,				lowerCAmelCase  :			Tuple=0.1			,				lowerCAmelCase  :			Any=0.1			,				lowerCAmelCase  :			Optional[Any]=5_12			,				lowerCAmelCase  :			Any=16			,				lowerCAmelCase  :			List[str]=2			,				lowerCAmelCase  :			Optional[Any]=0.02			,				lowerCAmelCase  :			Any=3			,				lowerCAmelCase  :			Any=4			,				lowerCAmelCase  :			Optional[int]=None			,				lowerCAmelCase  :			Optional[int]=0			,				)	->				Union[str, Any]:




         """simple docstring"""

         lowercase__     =  parent
         lowercase__     =  batch_size
         lowercase__     =  seq_length
         lowercase__     =  is_training
         lowercase__     =  use_input_mask
         lowercase__     =  use_token_type_ids
         lowercase__     =  use_labels
         lowercase__     =  vocab_size
         lowercase__     =  hidden_size
         lowercase__     =  num_hidden_layers
         lowercase__     =  num_attention_heads
         lowercase__     =  intermediate_size
         lowercase__     =  hidden_act
         lowercase__     =  hidden_dropout_prob
         lowercase__     =  attention_probs_dropout_prob
         lowercase__     =  max_position_embeddings
         lowercase__     =  type_vocab_size
         lowercase__     =  type_sequence_label_size
         lowercase__     =  initializer_range
         lowercase__     =  num_labels
         lowercase__     =  num_choices
         lowercase__     =  scope
         lowercase__     =  projection_dim


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




         """simple docstring"""

         lowercase__     =  ids_tensor([self.batch_size, self.seq_length]			,				self.vocab_size)

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

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

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

         lowercase__     =  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=lowerCAmelCase			,				initializer_range=self.initializer_range			,				)
         lowercase__     =  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  :			Dict			,				lowerCAmelCase  :			List[Any]			,				lowerCAmelCase  :			Optional[int]			,				lowerCAmelCase  :			Optional[int]			,				lowerCAmelCase  :			List[Any]			,				lowerCAmelCase  :			str			,				lowerCAmelCase  :			Any			,				lowerCAmelCase  :			Dict)	->				Any:




         """simple docstring"""

         lowercase__     =  TFDPRContextEncoder(config=lowerCAmelCase)
         lowercase__     =  model(lowerCAmelCase			,				attention_mask=lowerCAmelCase			,				token_type_ids=lowerCAmelCase)
         lowercase__     =  model(lowerCAmelCase			,				token_type_ids=lowerCAmelCase)
         lowercase__     =  model(lowerCAmelCase)
         self.parent.assertEqual(result.pooler_output.shape			,				(self.batch_size, self.projection_dim or self.hidden_size))


   def 			UpperCAmelCase		(						self  :			str			,				lowerCAmelCase  :			Any			,				lowerCAmelCase  :			int			,				lowerCAmelCase  :			Optional[Any]			,				lowerCAmelCase  :			Any			,				lowerCAmelCase  :			List[Any]			,				lowerCAmelCase  :			Optional[int]			,				lowerCAmelCase  :			Optional[int])	->				Optional[Any]:




         """simple docstring"""

         lowercase__     =  TFDPRQuestionEncoder(config=lowerCAmelCase)
         lowercase__     =  model(lowerCAmelCase			,				attention_mask=lowerCAmelCase			,				token_type_ids=lowerCAmelCase)
         lowercase__     =  model(lowerCAmelCase			,				token_type_ids=lowerCAmelCase)
         lowercase__     =  model(lowerCAmelCase)
         self.parent.assertEqual(result.pooler_output.shape			,				(self.batch_size, self.projection_dim or self.hidden_size))


   def 			UpperCAmelCase		(						self  :			str			,				lowerCAmelCase  :			Any			,				lowerCAmelCase  :			List[Any]			,				lowerCAmelCase  :			Dict			,				lowerCAmelCase  :			str			,				lowerCAmelCase  :			List[str]			,				lowerCAmelCase  :			List[str]			,				lowerCAmelCase  :			Any)	->				List[Any]:




         """simple docstring"""

         lowercase__     =  TFDPRReader(config=lowerCAmelCase)
         lowercase__     =  model(lowerCAmelCase			,				attention_mask=lowerCAmelCase)

         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  :			Tuple)	->				Optional[Any]:




         """simple docstring"""

         lowercase__     =  self.prepare_config_and_inputs()
         (
             (
             lowercase__	
         ),	(
             lowercase__	
         ),	(
             lowercase__	
         ),	(
             lowercase__	
         ),	(
             lowercase__	
         ),	(
             lowercase__	
         ),	(
             lowercase__	
         ),		
         )     =  config_and_inputs
         lowercase__     =  {'input_ids': input_ids}
         return config, inputs_dict



@require_tf
class       UpperCAmelCase__(							lowerCamelCase     ,							lowerCamelCase     ,							unittest.TestCase	):
   A    :  Optional[Any]       				=    (
       (
           TFDPRContextEncoder,
           TFDPRQuestionEncoder,
           TFDPRReader,
       )
       if is_tf_available()
       else ()
   )
   A    :  Optional[Any]       				=    {"feature-extraction": TFDPRQuestionEncoder} if is_tf_available() else {}

   A    :  Any       				=    False
   A    :  Union[str, Any]       				=    False
   A    :  Dict       				=    False
   A    :  Tuple       				=    False
   A    :  Any       				=    False


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




         """simple docstring"""

         lowercase__     =  TFDPRModelTester(self)
         lowercase__     =  ConfigTester(self			,				config_class=lowerCAmelCase			,				hidden_size=37)


   def 			UpperCAmelCase		(						self  :			Optional[int])	->				Dict:




         """simple docstring"""

         self.config_tester.run_common_tests()


   def 			UpperCAmelCase		(						self  :			str)	->				Dict:




         """simple docstring"""

         lowercase__     =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_dpr_context_encoder(*lowerCAmelCase)


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




         """simple docstring"""

         lowercase__     =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_dpr_question_encoder(*lowerCAmelCase)


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




         """simple docstring"""

         lowercase__     =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_dpr_reader(*lowerCAmelCase)



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




         """simple docstring"""

         for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               lowercase__     =  TFDPRContextEncoder.from_pretrained(lowerCAmelCase)
               self.assertIsNotNone(lowerCAmelCase)

         for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               lowercase__     =  TFDPRContextEncoder.from_pretrained(lowerCAmelCase)
               self.assertIsNotNone(lowerCAmelCase)

         for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               lowercase__     =  TFDPRQuestionEncoder.from_pretrained(lowerCAmelCase)
               self.assertIsNotNone(lowerCAmelCase)

         for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               lowercase__     =  TFDPRReader.from_pretrained(lowerCAmelCase)
               self.assertIsNotNone(lowerCAmelCase)



@require_tf
class       UpperCAmelCase__(							unittest.TestCase	):


   @slow
   def 			UpperCAmelCase		(						self  :			Any)	->				Any:




         """simple docstring"""

         lowercase__     =  TFDPRQuestionEncoder.from_pretrained('facebook/dpr-question_encoder-single-nq-base')

         lowercase__     =  tf.constant(
             [[1_01, 75_92, 10_10, 20_03, 20_26, 38_99, 1_01_40, 10_29, 1_02]])  # [CLS] hello, is my dog cute? [SEP]
         lowercase__     =  model(lowerCAmelCase)[0]  # embedding shape = (1, 768)
         # compare the actual values for a slice.
         lowercase__     =  tf.constant(
             [
                 [
                     0.03_23_62_53,
                     0.12_75_33_35,
                     0.16_81_85_09,
                     0.00_27_97_86,
                     0.3_89_69_33,
                     0.24_26_49_45,
                     0.2_17_89_71,
                     -0.02_33_52_27,
                     -0.08_48_19_59,
                     -0.14_32_41_17,
                 ]
             ])
         self.assertTrue(numpy.allclose(output[:, :10].numpy()			,				expected_slice.numpy()			,				atol=1E-4))





import torch

from diffusers import DDIMParallelScheduler

from .test_schedulers import SchedulerCommonTest



class       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




   A    :  str       				=    (DDIMParallelScheduler,)
   A    :  Any       				=    (("eta", 0.0), ("num_inference_steps", 50))


   def 			UpperCAmelCase		(						self  :			Union[str, Any]			,				**lowerCAmelCase  :			Optional[int])	->				Dict:




         """simple docstring"""

         lowercase__     =  {
             'num_train_timesteps': 10_00,
             'beta_start': 0.00_01,
             'beta_end': 0.02,
             'beta_schedule': 'linear',
             'clip_sample': True,
         }

         config.update(**lowerCAmelCase)
         return config


   def 			UpperCAmelCase		(						self  :			int			,				**lowerCAmelCase  :			str)	->				Union[str, Any]:




         """simple docstring"""

         lowercase__     =  self.scheduler_classes[0]
         lowercase__     =  self.get_scheduler_config(**lowerCAmelCase)
         lowercase__     =  scheduler_class(**lowerCAmelCase)

         lowercase__,	lowercase__     =  10, 0.0

         lowercase__     =  self.dummy_model()
         lowercase__     =  self.dummy_sample_deter

         scheduler.set_timesteps(lowerCAmelCase)

         for t in scheduler.timesteps:
               lowercase__     =  model(lowerCAmelCase			,				lowerCAmelCase)
               lowercase__     =  scheduler.step(lowerCAmelCase			,				lowerCAmelCase			,				lowerCAmelCase			,				lowerCAmelCase).prev_sample

         return sample


   def 			UpperCAmelCase		(						self  :			Tuple)	->				int:




         """simple docstring"""

         for timesteps in [1_00, 5_00, 10_00]:
               self.check_over_configs(num_train_timesteps=lowerCAmelCase)


   def 			UpperCAmelCase		(						self  :			Tuple)	->				Any:




         """simple docstring"""

         for steps_offset in [0, 1]:
               self.check_over_configs(steps_offset=lowerCAmelCase)

         lowercase__     =  self.scheduler_classes[0]
         lowercase__     =  self.get_scheduler_config(steps_offset=1)
         lowercase__     =  scheduler_class(**lowerCAmelCase)
         scheduler.set_timesteps(5)
         assert torch.equal(scheduler.timesteps			,				torch.LongTensor([8_01, 6_01, 4_01, 2_01, 1]))


   def 			UpperCAmelCase		(						self  :			str)	->				Tuple:




         """simple docstring"""

         for beta_start, beta_end in zip([0.00_01, 0.0_01, 0.01, 0.1]			,				[0.0_02, 0.02, 0.2, 2]):
               self.check_over_configs(beta_start=lowerCAmelCase			,				beta_end=lowerCAmelCase)


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




         """simple docstring"""

         for schedule in ["linear", "squaredcos_cap_v2"]:
               self.check_over_configs(beta_schedule=lowerCAmelCase)


   def 			UpperCAmelCase		(						self  :			List[str])	->				List[str]:




         """simple docstring"""

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


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




         """simple docstring"""

         for clip_sample in [True, False]:
               self.check_over_configs(clip_sample=lowerCAmelCase)


   def 			UpperCAmelCase		(						self  :			Optional[int])	->				int:




         """simple docstring"""

         for timestep_spacing in ["trailing", "leading"]:
               self.check_over_configs(timestep_spacing=lowerCAmelCase)


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




         """simple docstring"""

         for rescale_betas_zero_snr in [True, False]:
               self.check_over_configs(rescale_betas_zero_snr=lowerCAmelCase)


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




         """simple docstring"""

         self.check_over_configs(thresholding=lowerCAmelCase)
         for threshold in [0.5, 1.0, 2.0]:
               for prediction_type in ["epsilon", "v_prediction"]:
                     self.check_over_configs(
                         thresholding=lowerCAmelCase			,				prediction_type=lowerCAmelCase			,				sample_max_value=lowerCAmelCase			,				)


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




         """simple docstring"""

         for t in [1, 10, 49]:
               self.check_over_forward(time_step=lowerCAmelCase)


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




         """simple docstring"""

         for t, num_inference_steps in zip([1, 10, 50]			,				[10, 50, 5_00]):
               self.check_over_forward(time_step=lowerCAmelCase			,				num_inference_steps=lowerCAmelCase)


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




         """simple docstring"""

         for t, eta in zip([1, 10, 49]			,				[0.0, 0.5, 1.0]):
               self.check_over_forward(time_step=lowerCAmelCase			,				eta=lowerCAmelCase)


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




         """simple docstring"""

         lowercase__     =  self.scheduler_classes[0]
         lowercase__     =  self.get_scheduler_config()
         lowercase__     =  scheduler_class(**lowerCAmelCase)

         assert torch.sum(torch.abs(scheduler._get_variance(0			,				0) - 0.0)) < 1E-5
         assert torch.sum(torch.abs(scheduler._get_variance(4_20			,				4_00) - 0.1_47_71)) < 1E-5
         assert torch.sum(torch.abs(scheduler._get_variance(9_80			,				9_60) - 0.3_24_60)) < 1E-5
         assert torch.sum(torch.abs(scheduler._get_variance(0			,				0) - 0.0)) < 1E-5
         assert torch.sum(torch.abs(scheduler._get_variance(4_87			,				4_86) - 0.0_09_79)) < 1E-5
         assert torch.sum(torch.abs(scheduler._get_variance(9_99			,				9_98) - 0.02)) < 1E-5


   def 			UpperCAmelCase		(						self  :			Dict)	->				Tuple:




         """simple docstring"""

         lowercase__     =  self.scheduler_classes[0]
         lowercase__     =  self.get_scheduler_config()
         lowercase__     =  scheduler_class(**lowerCAmelCase)

         lowercase__,	lowercase__     =  10, 0.0
         scheduler.set_timesteps(lowerCAmelCase)

         lowercase__     =  self.dummy_model()
         lowercase__     =  self.dummy_sample_deter
         lowercase__     =  self.dummy_sample_deter + 0.1
         lowercase__     =  self.dummy_sample_deter - 0.1

         lowercase__     =  samplea.shape[0]
         lowercase__     =  torch.stack([samplea, samplea, samplea]			,				dim=0)
         lowercase__     =  torch.arange(lowerCAmelCase)[0:3, None].repeat(1			,				lowerCAmelCase)

         lowercase__     =  model(samples.flatten(0			,				1)			,				timesteps.flatten(0			,				1))
         lowercase__     =  scheduler.batch_step_no_noise(lowerCAmelCase			,				timesteps.flatten(0			,				1)			,				samples.flatten(0			,				1)			,				lowerCAmelCase)

         lowercase__     =  torch.sum(torch.abs(lowerCAmelCase))
         lowercase__     =  torch.mean(torch.abs(lowerCAmelCase))

         assert abs(result_sum.item() - 11_47.79_04) < 1E-2
         assert abs(result_mean.item() - 0.49_82) < 1E-3


   def 			UpperCAmelCase		(						self  :			Any)	->				int:




         """simple docstring"""

         lowercase__     =  self.full_loop()

         lowercase__     =  torch.sum(torch.abs(lowerCAmelCase))
         lowercase__     =  torch.mean(torch.abs(lowerCAmelCase))

         assert abs(result_sum.item() - 1_72.00_67) < 1E-2
         assert abs(result_mean.item() - 0.22_39_67) < 1E-3


   def 			UpperCAmelCase		(						self  :			int)	->				List[Any]:




         """simple docstring"""

         lowercase__     =  self.full_loop(prediction_type='v_prediction')

         lowercase__     =  torch.sum(torch.abs(lowerCAmelCase))
         lowercase__     =  torch.mean(torch.abs(lowerCAmelCase))

         assert abs(result_sum.item() - 52.53_02) < 1E-2
         assert abs(result_mean.item() - 0.06_84) < 1E-3


   def 			UpperCAmelCase		(						self  :			str)	->				Dict:




         """simple docstring"""

         lowercase__     =  self.full_loop(set_alpha_to_one=lowerCAmelCase			,				beta_start=0.01)
         lowercase__     =  torch.sum(torch.abs(lowerCAmelCase))
         lowercase__     =  torch.mean(torch.abs(lowerCAmelCase))

         assert abs(result_sum.item() - 1_49.82_95) < 1E-2
         assert abs(result_mean.item() - 0.19_51) < 1E-3



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




         """simple docstring"""

         lowercase__     =  self.full_loop(set_alpha_to_one=lowerCAmelCase			,				beta_start=0.01)
         lowercase__     =  torch.sum(torch.abs(lowerCAmelCase))
         lowercase__     =  torch.mean(torch.abs(lowerCAmelCase))

         assert abs(result_sum.item() - 1_49.07_84) < 1E-2
         assert abs(result_mean.item() - 0.19_41) < 1E-3
