


'''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    __magic_name__				:






       def __init__(     self    ,				snake_case    ,				snake_case=1_3    ,				snake_case=7    ,				snake_case=True    ,				snake_case=True    ,				snake_case=True    ,				snake_case=True    ,				snake_case=9_9    ,				snake_case=3_2    ,				snake_case=2    ,				snake_case=4    ,				snake_case=3_7    ,				snake_case="gelu"    ,				snake_case=0.1    ,				snake_case=0.1    ,				snake_case=5_1_2    ,				snake_case=1_6    ,				snake_case=2    ,				snake_case=0.02    ,				snake_case=3    ,				snake_case=4    ,				snake_case=None    ,				snake_case=0    ,				)		->    Any:




          '''simple docstring'''
          _UpperCAmelCase  :  str				       =parent
          _UpperCAmelCase  :  Optional[Any]				       =batch_size
          _UpperCAmelCase  :  Tuple				       =seq_length
          _UpperCAmelCase  :  Any				       =is_training
          _UpperCAmelCase  :  Optional[Any]				       =use_input_mask
          _UpperCAmelCase  :  str				       =use_token_type_ids
          _UpperCAmelCase  :  str				       =use_labels
          _UpperCAmelCase  :  str				       =vocab_size
          _UpperCAmelCase  :  Tuple				       =hidden_size
          _UpperCAmelCase  :  List[Any]				       =num_hidden_layers
          _UpperCAmelCase  :  int				       =num_attention_heads
          _UpperCAmelCase  :  Union[str, Any]				       =intermediate_size
          _UpperCAmelCase  :  Optional[Any]				       =hidden_act
          _UpperCAmelCase  :  List[Any]				       =hidden_dropout_prob
          _UpperCAmelCase  :  List[Any]				       =attention_probs_dropout_prob
          _UpperCAmelCase  :  Optional[int]				       =max_position_embeddings
          _UpperCAmelCase  :  Dict				       =type_vocab_size
          _UpperCAmelCase  :  Optional[int]				       =type_sequence_label_size
          _UpperCAmelCase  :  int				       =initializer_range
          _UpperCAmelCase  :  str				       =num_labels
          _UpperCAmelCase  :  Any				       =num_choices
          _UpperCAmelCase  :  Any				       =scope
          _UpperCAmelCase  :  Union[str, Any]				       =projection_dim






       def      lowerCAmelCase						(     self)		->    List[Any]:




          '''simple docstring'''
          _UpperCAmelCase  :  List[Any]				       =ids_tensor([self.batch_size, self.seq_length]    ,				self.vocab_size)

          _UpperCAmelCase  :  List[str]				       =None
          if self.use_input_mask:
             # follow test_modeling_tf_ctrl.py
             _UpperCAmelCase  :  List[Any]				       =random_attention_mask([self.batch_size, self.seq_length])

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

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

          _UpperCAmelCase  :  Union[str, Any]				       =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    ,				)
          _UpperCAmelCase  :  Any				       =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      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case)		->    Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase  :  List[str]				       =TFDPRContextEncoder(config=__a)
          _UpperCAmelCase  :  Union[str, Any]				       =model(__a    ,				attention_mask=__a    ,				token_type_ids=__a)
          _UpperCAmelCase  :  List[str]				       =model(__a    ,				token_type_ids=__a)
          _UpperCAmelCase  :  Dict				       =model(__a)
          self.parent.assertEqual(result.pooler_output.shape    ,				(self.batch_size, self.projection_dim or self.hidden_size))






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




          '''simple docstring'''
          _UpperCAmelCase  :  List[str]				       =TFDPRQuestionEncoder(config=__a)
          _UpperCAmelCase  :  str				       =model(__a    ,				attention_mask=__a    ,				token_type_ids=__a)
          _UpperCAmelCase  :  Union[str, Any]				       =model(__a    ,				token_type_ids=__a)
          _UpperCAmelCase  :  int				       =model(__a)
          self.parent.assertEqual(result.pooler_output.shape    ,				(self.batch_size, self.projection_dim or self.hidden_size))






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case)		->    Dict:




          '''simple docstring'''
          _UpperCAmelCase  :  Optional[Any]				       =TFDPRReader(config=__a)
          _UpperCAmelCase  :  str				       =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      lowerCAmelCase						(     self)		->    Any:




          '''simple docstring'''
          _UpperCAmelCase  :  List[str]				       =self.prepare_config_and_inputs()
          (
              (
              _UpperCAmelCase    
          )       ,			(
              _UpperCAmelCase    
          )       ,			(
              _UpperCAmelCase    
          )       ,			(
              _UpperCAmelCase    
          )       ,			(
              _UpperCAmelCase    
          )       ,			(
              _UpperCAmelCase    
          )       ,			(
              _UpperCAmelCase    
          )       ,			    
          )  :  Optional[int]				       =config_and_inputs
          _UpperCAmelCase  :  List[str]				       ={'input_ids': input_ids}
          return config, inputs_dict






@require_tf
class    __magic_name__				(						lowerCAmelCase   ,lowerCAmelCase   ,unittest.TestCase							):
       UpperCAmelCase		 =(
           (
               TFDPRContextEncoder,
               TFDPRQuestionEncoder,
               TFDPRReader,
           )
           if is_tf_available()
           else ()
       )
       UpperCAmelCase		 ={"feature-extraction": TFDPRQuestionEncoder} if is_tf_available() else {}

       UpperCAmelCase		 =False
       UpperCAmelCase		 =False
       UpperCAmelCase		 =False
       UpperCAmelCase		 =False
       UpperCAmelCase		 =False






       def      lowerCAmelCase						(     self)		->    Tuple:




          '''simple docstring'''
          _UpperCAmelCase  :  Dict				       =TFDPRModelTester(self)
          _UpperCAmelCase  :  Optional[int]				       =ConfigTester(self    ,				config_class=__a    ,				hidden_size=3_7)






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




          '''simple docstring'''
          self.config_tester.run_common_tests()






       def      lowerCAmelCase						(     self)		->    Tuple:




          '''simple docstring'''
          _UpperCAmelCase  :  Optional[Any]				       =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_dpr_context_encoder(*__a)






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




          '''simple docstring'''
          _UpperCAmelCase  :  Dict				       =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_dpr_question_encoder(*__a)






       def      lowerCAmelCase						(     self)		->    Dict:




          '''simple docstring'''
          _UpperCAmelCase  :  Dict				       =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_dpr_reader(*__a)






       @slow
       def      lowerCAmelCase						(     self)		->    str:




          '''simple docstring'''
          for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _UpperCAmelCase  :  Any				       =TFDPRContextEncoder.from_pretrained(__a)
             self.assertIsNotNone(__a)

          for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _UpperCAmelCase  :  int				       =TFDPRContextEncoder.from_pretrained(__a)
             self.assertIsNotNone(__a)

          for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _UpperCAmelCase  :  Union[str, Any]				       =TFDPRQuestionEncoder.from_pretrained(__a)
             self.assertIsNotNone(__a)

          for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _UpperCAmelCase  :  Tuple				       =TFDPRReader.from_pretrained(__a)
             self.assertIsNotNone(__a)






@require_tf
class    __magic_name__				(						unittest.TestCase							):






       @slow
       def      lowerCAmelCase						(     self)		->    Tuple:




          '''simple docstring'''
          _UpperCAmelCase  :  int				       =TFDPRQuestionEncoder.from_pretrained('facebook/dpr-question_encoder-single-nq-base')

          _UpperCAmelCase  :  Optional[int]				       =tf.constant(
              [[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_0_3, 2_0_2_6, 3_8_9_9, 1_0_1_4_0, 1_0_2_9, 1_0_2]])  # [CLS] hello, is my dog cute? [SEP]
          _UpperCAmelCase  :  Optional[int]				       =model(__a)[0]  # embedding shape = (1, 768)
          # compare the actual values for a slice.
          _UpperCAmelCase  :  Optional[Any]				       =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[:, :1_0].numpy()    ,				expected_slice.numpy()    ,				atol=1E-4))



"""simple docstring"""





import torch

from diffusers import DDPMScheduler

from .test_schedulers import SchedulerCommonTest






class        _UpperCAmelCase(       lowerCamelCase      ):
		lowercase__         = (DDPMScheduler,)



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






							'''simple docstring'''
							_UpperCamelCase									=    {
							    '''num_train_timesteps''': 10_00,
							    '''beta_start''': 0.0001,
							    '''beta_end''': 0.02,
							    '''beta_schedule''': '''linear''',
							    '''variance_type''': '''fixed_small''',
							    '''clip_sample''': True,
							}

							config.update(**__a)
							return config



		def 		UpperCAmelCase				(   self)			->			Dict:






							'''simple docstring'''
							for timesteps in [1, 5, 1_00, 10_00]:
												self.check_over_configs(num_train_timesteps=__a)



		def 		UpperCAmelCase				(   self)			->			int:






							'''simple docstring'''
							for beta_start, beta_end in zip([0.0001, 0.001, 0.01, 0.1]							,						[0.002, 0.02, 0.2, 2]):
												self.check_over_configs(beta_start=__a							,						beta_end=__a)



		def 		UpperCAmelCase				(   self)			->			Tuple:






							'''simple docstring'''
							for schedule in ["linear", "squaredcos_cap_v2"]:
												self.check_over_configs(beta_schedule=__a)



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






							'''simple docstring'''
							for variance in ["fixed_small", "fixed_large", "other"]:
												self.check_over_configs(variance_type=__a)



		def 		UpperCAmelCase				(   self)			->			Dict:






							'''simple docstring'''
							for clip_sample in [True, False]:
												self.check_over_configs(clip_sample=__a)



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






							'''simple docstring'''
							self.check_over_configs(thresholding=__a)
							for threshold in [0.5, 1.0, 2.0]:
												for prediction_type in ["epsilon", "sample", "v_prediction"]:
																	self.check_over_configs(
																	    thresholding=__a							,						prediction_type=__a							,						sample_max_value=__a							,						)



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






							'''simple docstring'''
							for prediction_type in ["epsilon", "sample", "v_prediction"]:
												self.check_over_configs(prediction_type=__a)



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






							'''simple docstring'''
							for t in [0, 5_00, 9_99]:
												self.check_over_forward(time_step=__a)



		def 		UpperCAmelCase				(   self)			->			Tuple:






							'''simple docstring'''
							_UpperCamelCase									=    self.scheduler_classes[0]
							_UpperCamelCase									=    self.get_scheduler_config()
							_UpperCamelCase									=    scheduler_class(**__a)

							assert torch.sum(torch.abs(scheduler._get_variance(0) - 0.0)) < 1e-5
							assert torch.sum(torch.abs(scheduler._get_variance(4_87) - 0.0_0979)) < 1e-5
							assert torch.sum(torch.abs(scheduler._get_variance(9_99) - 0.02)) < 1e-5



		def 		UpperCAmelCase				(   self)			->			str:






							'''simple docstring'''
							_UpperCamelCase									=    self.scheduler_classes[0]
							_UpperCamelCase									=    self.get_scheduler_config()
							_UpperCamelCase									=    scheduler_class(**__a)

							_UpperCamelCase									=    len(__a)

							_UpperCamelCase									=    self.dummy_model()
							_UpperCamelCase									=    self.dummy_sample_deter
							_UpperCamelCase									=    torch.manual_seed(0)

							for t in reversed(range(__a)):
												# 1. predict noise residual
												_UpperCamelCase									=    model(__a							,						__a)

												# 2. predict previous mean of sample x_t-1
												_UpperCamelCase									=    scheduler.step(__a							,						__a							,						__a							,						generator=__a).prev_sample

												# if t > 0:
												#     noise = self.dummy_sample_deter
												#     variance = scheduler.get_variance(t) ** (0.5) * noise
												#
												# sample = pred_prev_sample + variance
												_UpperCamelCase									=    pred_prev_sample

							_UpperCamelCase									=    torch.sum(torch.abs(__a))
							_UpperCamelCase									=    torch.mean(torch.abs(__a))

							assert abs(result_sum.item() - 258.9606) < 1e-2
							assert abs(result_mean.item() - 0.3372) < 1e-3



		def 		UpperCAmelCase				(   self)			->			str:






							'''simple docstring'''
							_UpperCamelCase									=    self.scheduler_classes[0]
							_UpperCamelCase									=    self.get_scheduler_config(prediction_type='''v_prediction''')
							_UpperCamelCase									=    scheduler_class(**__a)

							_UpperCamelCase									=    len(__a)

							_UpperCamelCase									=    self.dummy_model()
							_UpperCamelCase									=    self.dummy_sample_deter
							_UpperCamelCase									=    torch.manual_seed(0)

							for t in reversed(range(__a)):
												# 1. predict noise residual
												_UpperCamelCase									=    model(__a							,						__a)

												# 2. predict previous mean of sample x_t-1
												_UpperCamelCase									=    scheduler.step(__a							,						__a							,						__a							,						generator=__a).prev_sample

												# if t > 0:
												#     noise = self.dummy_sample_deter
												#     variance = scheduler.get_variance(t) ** (0.5) * noise
												#
												# sample = pred_prev_sample + variance
												_UpperCamelCase									=    pred_prev_sample

							_UpperCamelCase									=    torch.sum(torch.abs(__a))
							_UpperCamelCase									=    torch.mean(torch.abs(__a))

							assert abs(result_sum.item() - 202.0296) < 1e-2
							assert abs(result_mean.item() - 0.2631) < 1e-3



		def 		UpperCAmelCase				(   self)			->			Any:






							'''simple docstring'''
							_UpperCamelCase									=    self.scheduler_classes[0]
							_UpperCamelCase									=    self.get_scheduler_config()
							_UpperCamelCase									=    scheduler_class(**__a)

							_UpperCamelCase									=    [1_00, 87, 50, 1, 0]

							scheduler.set_timesteps(timesteps=__a)

							_UpperCamelCase									=    scheduler.timesteps

							for i, timestep in enumerate(__a):
												if i == len(__a) - 1:
																	_UpperCamelCase									=    -1
												else:
																	_UpperCamelCase									=    timesteps[i + 1]

												_UpperCamelCase									=    scheduler.previous_timestep(__a)
												_UpperCamelCase									=    prev_t.item()

												self.assertEqual(__a							,						__a)



		def 		UpperCAmelCase				(   self)			->			Dict:






							'''simple docstring'''
							_UpperCamelCase									=    self.scheduler_classes[0]
							_UpperCamelCase									=    self.get_scheduler_config()
							_UpperCamelCase									=    scheduler_class(**__a)

							_UpperCamelCase									=    [1_00, 87, 50, 51, 0]

							with self.assertRaises(__a							,						msg='''`custom_timesteps` must be in descending order.'''):
												scheduler.set_timesteps(timesteps=__a)



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






							'''simple docstring'''
							_UpperCamelCase									=    self.scheduler_classes[0]
							_UpperCamelCase									=    self.get_scheduler_config()
							_UpperCamelCase									=    scheduler_class(**__a)

							_UpperCamelCase									=    [1_00, 87, 50, 1, 0]
							_UpperCamelCase									=    len(__a)

							with self.assertRaises(__a							,						msg='''Can only pass one of `num_inference_steps` or `custom_timesteps`.'''):
												scheduler.set_timesteps(num_inference_steps=__a							,						timesteps=__a)




		def 		UpperCAmelCase				(   self)			->			Tuple:






							'''simple docstring'''
							_UpperCamelCase									=    self.scheduler_classes[0]
							_UpperCamelCase									=    self.get_scheduler_config()
							_UpperCamelCase									=    scheduler_class(**__a)

							_UpperCamelCase									=    [scheduler.config.num_train_timesteps]

							with self.assertRaises(
							    __a							,						msg='''`timesteps` must start before `self.config.train_timesteps`: {scheduler.config.num_train_timesteps}}'''							,						):
												scheduler.set_timesteps(timesteps=__a)






