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


   def __init__(self	:		Optional[Any]			,    __UpperCAmelCase	:		str			,    __UpperCAmelCase	:		int=1_3			,    __UpperCAmelCase	:		Union[str, Any]=7			,    __UpperCAmelCase	:		int=True			,    __UpperCAmelCase	:		Dict=True			,    __UpperCAmelCase	:		Tuple=True			,    __UpperCAmelCase	:		Any=True			,    __UpperCAmelCase	:		List[str]=9_9			,    __UpperCAmelCase	:		Optional[Any]=3_2			,    __UpperCAmelCase	:		Dict=2			,    __UpperCAmelCase	:		List[str]=4			,    __UpperCAmelCase	:		List[str]=3_7			,    __UpperCAmelCase	:		Dict="gelu"			,    __UpperCAmelCase	:		Dict=0.1			,    __UpperCAmelCase	:		Any=0.1			,    __UpperCAmelCase	:		str=5_1_2			,    __UpperCAmelCase	:		Union[str, Any]=1_6			,    __UpperCAmelCase	:		Optional[Any]=2			,    __UpperCAmelCase	:		Any=0.02			,    __UpperCAmelCase	:		List[str]=3			,    __UpperCAmelCase	:		List[str]=4			,    __UpperCAmelCase	:		Optional[int]=None			,    __UpperCAmelCase	:		List[str]=0			,    )		->							Dict:
         """simple docstring"""



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


   def      lowercase_ (self	:		List[Any]						)		->							Optional[int]:
         """simple docstring"""



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

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

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

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

         UpperCAmelCase__								=    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			,    )
         UpperCAmelCase__								=    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      lowercase_ (self	:		Union[str, Any]			,    __UpperCAmelCase	:		Union[str, Any]			,    __UpperCAmelCase	:		Optional[int]			,    __UpperCAmelCase	:		str			,    __UpperCAmelCase	:		List[Any]			,    __UpperCAmelCase	:		int			,    __UpperCAmelCase	:		Tuple			,    __UpperCAmelCase	:		Dict						)		->							Dict:
         """simple docstring"""



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


   def      lowercase_ (self	:		Tuple			,    __UpperCAmelCase	:		Optional[int]			,    __UpperCAmelCase	:		int			,    __UpperCAmelCase	:		Union[str, Any]			,    __UpperCAmelCase	:		Dict			,    __UpperCAmelCase	:		Optional[int]			,    __UpperCAmelCase	:		Tuple			,    __UpperCAmelCase	:		List[Any]						)		->							Optional[int]:
         """simple docstring"""



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


   def      lowercase_ (self	:		List[str]			,    __UpperCAmelCase	:		Tuple			,    __UpperCAmelCase	:		List[Any]			,    __UpperCAmelCase	:		Union[str, Any]			,    __UpperCAmelCase	:		Tuple			,    __UpperCAmelCase	:		Optional[Any]			,    __UpperCAmelCase	:		Optional[int]			,    __UpperCAmelCase	:		List[Any]						)		->							str:
         """simple docstring"""



         UpperCAmelCase__								=    TFDPRReader(config=lowerCAmelCase__						)
         UpperCAmelCase__								=    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      lowercase_ (self	:		str						)		->							Tuple:
         """simple docstring"""



         UpperCAmelCase__								=    self.prepare_config_and_inputs()
         (
             UpperCAmelCase__			
         )								=    config_and_inputs
         UpperCAmelCase__								=    {"input_ids": input_ids}
         return config, inputs_dict





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

   __UpperCAmelCase					:					int    =     False
   __UpperCAmelCase					:					Optional[int]    =     False
   __UpperCAmelCase					:					str    =     False
   __UpperCAmelCase					:					Any    =     False
   __UpperCAmelCase					:					Optional[int]    =     False


   def      lowercase_ (self	:		Optional[Any]						)		->							Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__								=    TFDPRModelTester(self						)
         UpperCAmelCase__								=    ConfigTester(self			,    config_class=lowerCAmelCase__			,    hidden_size=3_7						)


   def      lowercase_ (self	:		Optional[int]						)		->							str:
         """simple docstring"""



         self.config_tester.run_common_tests()


   def      lowercase_ (self	:		Optional[int]						)		->							List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_dpr_context_encoder(*lowerCAmelCase__						)


   def      lowercase_ (self	:		int						)		->							Optional[Any]:
         """simple docstring"""



         UpperCAmelCase__								=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_dpr_question_encoder(*lowerCAmelCase__						)


   def      lowercase_ (self	:		List[str]						)		->							List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_dpr_reader(*lowerCAmelCase__						)





   @slow
   def      lowercase_ (self	:		Union[str, Any]						)		->							Any:
         """simple docstring"""



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

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

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

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





@require_tf
class 	A				(							unittest.TestCase	):


   @slow
   def      lowercase_ (self	:		List[str]						)		->							List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=    TFDPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base"						)

         UpperCAmelCase__								=    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__								=    model(lowerCAmelCase__						)[0]  # embedding shape = (1, 768)
         # compare the actual values for a slice.
         UpperCAmelCase__								=    tf.constant(
             [
                 [
                     0.03236253,
                     0.12753335,
                     0.16818509,
                     0.00279786,
                     0.3896933,
                     0.24264945,
                     0.2178971,
                     -0.02335227,
                     -0.08481959,
                     -0.14324117,
                 ]
             ]						)
         self.assertTrue(numpy.allclose(output[:, :1_0].numpy()			,    expected_slice.numpy()			,    atol=1E-4						)						)










"""simple docstring"""




from ..utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_scipy_available,
    is_torch_available,
    is_torchsde_available,
)


try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ..utils.dummy_pt_objects import *  # noqa F403
else:
		from .scheduling_consistency_models import CMStochasticIterativeScheduler
		from .scheduling_ddim import DDIMScheduler
		from .scheduling_ddim_inverse import DDIMInverseScheduler
		from .scheduling_ddim_parallel import DDIMParallelScheduler
		from .scheduling_ddpm import DDPMScheduler
		from .scheduling_ddpm_parallel import DDPMParallelScheduler
		from .scheduling_deis_multistep import DEISMultistepScheduler
		from .scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler
		from .scheduling_dpmsolver_multistep_inverse import DPMSolverMultistepInverseScheduler
		from .scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
		from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler
		from .scheduling_euler_discrete import EulerDiscreteScheduler
		from .scheduling_heun_discrete import HeunDiscreteScheduler
		from .scheduling_ipndm import IPNDMScheduler
		from .scheduling_k_dpm_2_ancestral_discrete import KDPMaAncestralDiscreteScheduler
		from .scheduling_k_dpm_2_discrete import KDPMaDiscreteScheduler
		from .scheduling_karras_ve import KarrasVeScheduler
		from .scheduling_pndm import PNDMScheduler
		from .scheduling_repaint import RePaintScheduler
		from .scheduling_sde_ve import ScoreSdeVeScheduler
		from .scheduling_sde_vp import ScoreSdeVpScheduler
		from .scheduling_unclip import UnCLIPScheduler
		from .scheduling_unipc_multistep import UniPCMultistepScheduler
		from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
		from .scheduling_vq_diffusion import VQDiffusionScheduler

try:
		if not is_flax_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ..utils.dummy_flax_objects import *  # noqa F403
else:
		from .scheduling_ddim_flax import FlaxDDIMScheduler
		from .scheduling_ddpm_flax import FlaxDDPMScheduler
		from .scheduling_dpmsolver_multistep_flax import FlaxDPMSolverMultistepScheduler
		from .scheduling_karras_ve_flax import FlaxKarrasVeScheduler
		from .scheduling_lms_discrete_flax import FlaxLMSDiscreteScheduler
		from .scheduling_pndm_flax import FlaxPNDMScheduler
		from .scheduling_sde_ve_flax import FlaxScoreSdeVeScheduler
		from .scheduling_utils_flax import (
		    FlaxKarrasDiffusionSchedulers,
		    FlaxSchedulerMixin,
		    FlaxSchedulerOutput,
		    broadcast_to_shape_from_left,
		)


try:
		if not (is_torch_available() and is_scipy_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ..utils.dummy_torch_and_scipy_objects import *  # noqa F403
else:
		from .scheduling_lms_discrete import LMSDiscreteScheduler

try:
		if not (is_torch_available() and is_torchsde_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ..utils.dummy_torch_and_torchsde_objects import *  # noqa F403
else:
		from .scheduling_dpmsolver_sde import DPMSolverSDEScheduler


