



import tempfile

import torch

from diffusers import PNDMScheduler

from .test_schedulers import SchedulerCommonTest






class 						__a (     SCREAMING_SNAKE_CASE				):
 SCREAMING_SNAKE_CASE       =       (PNDMScheduler,)
 SCREAMING_SNAKE_CASE       =       (('num_inference_steps', 50),)



 def 		UpperCamelCase							(						self							:    Any				,	**snake_case_							:    Any)->							Optional[Any]:
      __lowerCAmelCase	   ={
          """num_train_timesteps""": 10_00,
          """beta_start""": 0.0_0_0_1,
          """beta_end""": 0.0_2,
          """beta_schedule""": """linear""",
      }

      config.update(**snake_case_)
      return config



 def 		UpperCamelCase							(						self							:    Union[str, Any]				,	snake_case_							:    List[Any]=0				,	**snake_case_							:    int)->							Union[str, Any]:
      __lowerCAmelCase	   =dict(self.forward_default_kwargs)
      __lowerCAmelCase	   =kwargs.pop("""num_inference_steps"""				,	snake_case_)
      __lowerCAmelCase	   =self.dummy_sample
      __lowerCAmelCase	   =0.1 * sample
      __lowerCAmelCase	   =[residual + 0.2, residual + 0.1_5, residual + 0.1, residual + 0.0_5]

      for scheduler_class in self.scheduler_classes:
           __lowerCAmelCase	   =self.get_scheduler_config(**snake_case_)
           __lowerCAmelCase	   =scheduler_class(**snake_case_)
           scheduler.set_timesteps(snake_case_)
           # copy over dummy past residuals
           __lowerCAmelCase	   =dummy_past_residuals[:]

           with tempfile.TemporaryDirectory() as tmpdirname:
                scheduler.save_config(snake_case_)
                __lowerCAmelCase	   =scheduler_class.from_pretrained(snake_case_)
                new_scheduler.set_timesteps(snake_case_)
                # copy over dummy past residuals
                __lowerCAmelCase	   =dummy_past_residuals[:]

           __lowerCAmelCase	   =scheduler.step_prk(snake_case_				,	snake_case_				,	snake_case_				,	**snake_case_).prev_sample
           __lowerCAmelCase	   =new_scheduler.step_prk(snake_case_				,	snake_case_				,	snake_case_				,	**snake_case_).prev_sample

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

           __lowerCAmelCase	   =scheduler.step_plms(snake_case_				,	snake_case_				,	snake_case_				,	**snake_case_).prev_sample
           __lowerCAmelCase	   =new_scheduler.step_plms(snake_case_				,	snake_case_				,	snake_case_				,	**snake_case_).prev_sample

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



 def 		UpperCamelCase							(						self							:    Union[str, Any])->							Optional[Any]:
      pass



 def 		UpperCamelCase							(						self							:    str				,	snake_case_							:    List[Any]=0				,	**snake_case_							:    List[Any])->							str:
      __lowerCAmelCase	   =dict(self.forward_default_kwargs)
      __lowerCAmelCase	   =kwargs.pop("""num_inference_steps"""				,	snake_case_)
      __lowerCAmelCase	   =self.dummy_sample
      __lowerCAmelCase	   =0.1 * sample
      __lowerCAmelCase	   =[residual + 0.2, residual + 0.1_5, residual + 0.1, residual + 0.0_5]

      for scheduler_class in self.scheduler_classes:
           __lowerCAmelCase	   =self.get_scheduler_config()
           __lowerCAmelCase	   =scheduler_class(**snake_case_)
           scheduler.set_timesteps(snake_case_)

           # copy over dummy past residuals (must be after setting timesteps)
           __lowerCAmelCase	   =dummy_past_residuals[:]

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

                # copy over dummy past residual (must be after setting timesteps)
                __lowerCAmelCase	   =dummy_past_residuals[:]

           __lowerCAmelCase	   =scheduler.step_prk(snake_case_				,	snake_case_				,	snake_case_				,	**snake_case_).prev_sample
           __lowerCAmelCase	   =new_scheduler.step_prk(snake_case_				,	snake_case_				,	snake_case_				,	**snake_case_).prev_sample

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

           __lowerCAmelCase	   =scheduler.step_plms(snake_case_				,	snake_case_				,	snake_case_				,	**snake_case_).prev_sample
           __lowerCAmelCase	   =new_scheduler.step_plms(snake_case_				,	snake_case_				,	snake_case_				,	**snake_case_).prev_sample

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



 def 		UpperCamelCase							(						self							:    Optional[Any]				,	**snake_case_							:    Dict)->							Tuple:
      __lowerCAmelCase	   =self.scheduler_classes[0]
      __lowerCAmelCase	   =self.get_scheduler_config(**snake_case_)
      __lowerCAmelCase	   =scheduler_class(**snake_case_)

      __lowerCAmelCase	   =10
      __lowerCAmelCase	   =self.dummy_model()
      __lowerCAmelCase	   =self.dummy_sample_deter
      scheduler.set_timesteps(snake_case_)

      for i, t in enumerate(scheduler.prk_timesteps):
           __lowerCAmelCase	   =model(snake_case_				,	snake_case_)
           __lowerCAmelCase	   =scheduler.step_prk(snake_case_				,	snake_case_				,	snake_case_).prev_sample

      for i, t in enumerate(scheduler.plms_timesteps):
           __lowerCAmelCase	   =model(snake_case_				,	snake_case_)
           __lowerCAmelCase	   =scheduler.step_plms(snake_case_				,	snake_case_				,	snake_case_).prev_sample

      return sample



 def 		UpperCamelCase							(						self							:    List[Any])->							Tuple:
      __lowerCAmelCase	   =dict(self.forward_default_kwargs)

      __lowerCAmelCase	   =kwargs.pop("""num_inference_steps"""				,	snake_case_)

      for scheduler_class in self.scheduler_classes:
           __lowerCAmelCase	   =self.get_scheduler_config()
           __lowerCAmelCase	   =scheduler_class(**snake_case_)

           __lowerCAmelCase	   =self.dummy_sample
           __lowerCAmelCase	   =0.1 * sample

           if num_inference_steps is not None and hasattr(snake_case_				,	"""set_timesteps"""):
                scheduler.set_timesteps(snake_case_)
           elif num_inference_steps is not None and not hasattr(snake_case_				,	"""set_timesteps"""):
                __lowerCAmelCase	   =num_inference_steps

           # copy over dummy past residuals (must be done after set_timesteps)
           __lowerCAmelCase	   =[residual + 0.2, residual + 0.1_5, residual + 0.1, residual + 0.0_5]
           __lowerCAmelCase	   =dummy_past_residuals[:]

           __lowerCAmelCase	   =scheduler.step_prk(snake_case_				,	0				,	snake_case_				,	**snake_case_).prev_sample
           __lowerCAmelCase	   =scheduler.step_prk(snake_case_				,	1				,	snake_case_				,	**snake_case_).prev_sample

           self.assertEqual(output_a.shape				,	sample.shape)
           self.assertEqual(output_a.shape				,	output_a.shape)

           __lowerCAmelCase	   =scheduler.step_plms(snake_case_				,	0				,	snake_case_				,	**snake_case_).prev_sample
           __lowerCAmelCase	   =scheduler.step_plms(snake_case_				,	1				,	snake_case_				,	**snake_case_).prev_sample

           self.assertEqual(output_a.shape				,	sample.shape)
           self.assertEqual(output_a.shape				,	output_a.shape)



 def 		UpperCamelCase							(						self							:    str)->							Optional[Any]:







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



 def 		UpperCamelCase							(						self							:    Dict)->							int:







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

      __lowerCAmelCase	   =self.scheduler_classes[0]
      __lowerCAmelCase	   =self.get_scheduler_config(steps_offset=1)
      __lowerCAmelCase	   =scheduler_class(**snake_case_)
      scheduler.set_timesteps(10)
      assert torch.equal(
          scheduler.timesteps				,	torch.LongTensor(
              [9_01, 8_51, 8_51, 8_01, 8_01, 7_51, 7_51, 7_01, 7_01, 6_51, 6_51, 6_01, 6_01, 5_01, 4_01, 3_01, 2_01, 1_01, 1])				,	)



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







      for beta_start, beta_end in zip([0.0_0_0_1, 0.0_0_1]				,	[0.0_0_2, 0.0_2]):
           self.check_over_configs(beta_start=snake_case_				,	beta_end=snake_case_)



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







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



 def 		UpperCamelCase							(						self							:    int)->							Tuple:







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



 def 		UpperCamelCase							(						self							:    Any)->							Any:







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



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







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



 def 		UpperCamelCase							(						self							:    int)->							Tuple:
      __lowerCAmelCase	   =27

      for scheduler_class in self.scheduler_classes:
           __lowerCAmelCase	   =self.dummy_sample
           __lowerCAmelCase	   =0.1 * sample

           __lowerCAmelCase	   =self.get_scheduler_config()
           __lowerCAmelCase	   =scheduler_class(**snake_case_)

           scheduler.set_timesteps(snake_case_)

           # before power of 3 fix, would error on first step, so we only need to do two
           for i, t in enumerate(scheduler.prk_timesteps[:2]):
                __lowerCAmelCase	   =scheduler.step_prk(snake_case_				,	snake_case_				,	snake_case_).prev_sample



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







      with self.assertRaises(snake_case_):
           __lowerCAmelCase	   =self.scheduler_classes[0]
           __lowerCAmelCase	   =self.get_scheduler_config()
           __lowerCAmelCase	   =scheduler_class(**snake_case_)

           scheduler.step_plms(self.dummy_sample				,	1				,	self.dummy_sample).prev_sample



 def 		UpperCamelCase							(						self							:    Any)->							Any:
      __lowerCAmelCase	   =self.full_loop()
      __lowerCAmelCase	   =torch.sum(torch.abs(snake_case_))
      __lowerCAmelCase	   =torch.mean(torch.abs(snake_case_))

      assert abs(result_sum.item() - 1_9_8.1_3_1_8) < 1e-2
      assert abs(result_mean.item() - 0.2_5_8_0) < 1e-3



 def 		UpperCamelCase							(						self							:    str)->							Optional[int]:
      __lowerCAmelCase	   =self.full_loop(prediction_type="""v_prediction""")
      __lowerCAmelCase	   =torch.sum(torch.abs(snake_case_))
      __lowerCAmelCase	   =torch.mean(torch.abs(snake_case_))

      assert abs(result_sum.item() - 6_7.3_9_8_6) < 1e-2
      assert abs(result_mean.item() - 0.0_8_7_8) < 1e-3



 def 		UpperCamelCase							(						self							:    List[Any])->							Optional[Any]:
      __lowerCAmelCase	   =self.full_loop(set_alpha_to_one=snake_case_				,	beta_start=0.0_1)
      __lowerCAmelCase	   =torch.sum(torch.abs(snake_case_))
      __lowerCAmelCase	   =torch.mean(torch.abs(snake_case_))

      assert abs(result_sum.item() - 2_3_0.0_3_9_9) < 1e-2
      assert abs(result_mean.item() - 0.2_9_9_5) < 1e-3



 def 		UpperCamelCase							(						self							:    Any)->							Union[str, Any]:
      __lowerCAmelCase	   =self.full_loop(set_alpha_to_one=snake_case_				,	beta_start=0.0_1)
      __lowerCAmelCase	   =torch.sum(torch.abs(snake_case_))
      __lowerCAmelCase	   =torch.mean(torch.abs(snake_case_))

      assert abs(result_sum.item() - 1_8_6.9_4_8_2) < 1e-2
      assert abs(result_mean.item() - 0.2_4_3_4) < 1e-3









from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase_	       =			logging.get_logger(__name__)

lowercase_	       =			{
    '''google/realm-cc-news-pretrained-embedder''': (
        '''https://huggingface.co/google/realm-cc-news-pretrained-embedder/resolve/main/config.json'''
    ),
    '''google/realm-cc-news-pretrained-encoder''': (
        '''https://huggingface.co/google/realm-cc-news-pretrained-encoder/resolve/main/config.json'''
    ),
    '''google/realm-cc-news-pretrained-scorer''': (
        '''https://huggingface.co/google/realm-cc-news-pretrained-scorer/resolve/main/config.json'''
    ),
    '''google/realm-cc-news-pretrained-openqa''': (
        '''https://huggingface.co/google/realm-cc-news-pretrained-openqa/aresolve/main/config.json'''
    ),
    '''google/realm-orqa-nq-openqa''': '''https://huggingface.co/google/realm-orqa-nq-openqa/resolve/main/config.json''',
    '''google/realm-orqa-nq-reader''': '''https://huggingface.co/google/realm-orqa-nq-reader/resolve/main/config.json''',
    '''google/realm-orqa-wq-openqa''': '''https://huggingface.co/google/realm-orqa-wq-openqa/resolve/main/config.json''',
    '''google/realm-orqa-wq-reader''': '''https://huggingface.co/google/realm-orqa-wq-reader/resolve/main/config.json''',
    # See all REALM models at https://huggingface.co/models?filter=realm
}






class 						__a (     SCREAMING_SNAKE_CASE				):
	SCREAMING_SNAKE_CASE       =       "realm"



	def __init__(						self							:    Tuple				,	snake_case_							:    Tuple=3_05_22				,	snake_case_							:    List[Any]=7_68				,	snake_case_							:    Any=1_28				,	snake_case_							:    Union[str, Any]=12				,	snake_case_							:    List[str]=12				,	snake_case_							:    Any=8				,	snake_case_							:    Any=30_72				,	snake_case_							:    int="gelu_new"				,	snake_case_							:    Optional[int]=0.1				,	snake_case_							:    Tuple=0.1				,	snake_case_							:    List[Any]=5_12				,	snake_case_							:    List[Any]=2				,	snake_case_							:    Any=0.0_2				,	snake_case_							:    List[str]=1e-12				,	snake_case_							:    Optional[Any]=2_56				,	snake_case_							:    Dict=10				,	snake_case_							:    str=1e-3				,	snake_case_							:    List[str]=5				,	snake_case_							:    Optional[Any]=3_20				,	snake_case_							:    Optional[Any]=13_35_37_18				,	snake_case_							:    Tuple=50_00				,	snake_case_							:    Union[str, Any]=1				,	snake_case_							:    Union[str, Any]=0				,	snake_case_							:    List[str]=2				,	**snake_case_							:    List[str]				,	)->							str:
						super().__init__(pad_token_id=snake_case_				,	bos_token_id=snake_case_				,	eos_token_id=snake_case_				,	**snake_case_)

						# Common config
						__lowerCAmelCase	   =vocab_size
						__lowerCAmelCase	   =max_position_embeddings
						__lowerCAmelCase	   =hidden_size
						__lowerCAmelCase	   =retriever_proj_size
						__lowerCAmelCase	   =num_hidden_layers
						__lowerCAmelCase	   =num_attention_heads
						__lowerCAmelCase	   =num_candidates
						__lowerCAmelCase	   =intermediate_size
						__lowerCAmelCase	   =hidden_act
						__lowerCAmelCase	   =hidden_dropout_prob
						__lowerCAmelCase	   =attention_probs_dropout_prob
						__lowerCAmelCase	   =initializer_range
						__lowerCAmelCase	   =type_vocab_size
						__lowerCAmelCase	   =layer_norm_eps

						# Reader config
						__lowerCAmelCase	   =span_hidden_size
						__lowerCAmelCase	   =max_span_width
						__lowerCAmelCase	   =reader_layer_norm_eps
						__lowerCAmelCase	   =reader_beam_size
						__lowerCAmelCase	   =reader_seq_len

						# Retrieval config
						__lowerCAmelCase	   =num_block_records
						__lowerCAmelCase	   =searcher_beam_size




