



'''simple docstring'''

import os
import tempfile
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
							import torch
							from torch import nn

							from transformers import (
							    Adafactor,
							    AdamW,
							    get_constant_schedule,
							    get_constant_schedule_with_warmup,
							    get_cosine_schedule_with_warmup,
							    get_cosine_with_hard_restarts_schedule_with_warmup,
							    get_inverse_sqrt_schedule,
							    get_linear_schedule_with_warmup,
							    get_polynomial_decay_schedule_with_warmup,
							)



def      __lowerCamelCase     (      lowerCAmelCase_						,			lowerCAmelCase_=10  )			->      str:
		_a     :      str        =							[]
		for _ in range(lowerCAmelCase__  ):
				lrs.append(scheduler.get_lr()[0]  )
				scheduler.step()
		return lrs



def      __lowerCamelCase     (      lowerCAmelCase_						,			lowerCAmelCase_=10  )			->      Any:
		_a     :      str        =							[]
		for step in range(lowerCAmelCase__  ):
				lrs.append(scheduler.get_lr()[0]  )
				scheduler.step()
				if step == num_steps // 2:
						with tempfile.TemporaryDirectory() as tmpdirname:
								_a     :      List[str]        =							os.path.join(lowerCAmelCase__						,			'schedule.bin'  )
								torch.save(scheduler.state_dict()						,			lowerCAmelCase__  )

								_a     :      List[str]        =							torch.load(lowerCAmelCase__  )
								scheduler.load_state_dict(lowerCAmelCase__  )
		return lrs

@require_torch
class 	__magic_name__						(    unittest.TestCase   ):







							def 							__lowercase			(							self					:       str	,_UpperCAmelCase					:       List[str]	,_UpperCAmelCase					:       str	,_UpperCAmelCase					:       Union[str, Any]		):
									self.assertEqual(len(__lowerCAmelCase		)	,len(__lowerCAmelCase		)		)
									for a, b in zip(__lowerCAmelCase	,__lowerCAmelCase		):
											self.assertAlmostEqual(__lowerCAmelCase	,__lowerCAmelCase	,delta=__lowerCAmelCase		)







							def 							__lowercase			(							self					:       Tuple		):
									_a     :      Any        =							torch.tensor([0.1, -0.2, -0.1]	,requires_grad=__lowerCAmelCase		)
									_a     :      Any        =							torch.tensor([0.4, 0.2, -0.5]		)
									_a     :      int        =							nn.MSELoss()
									# No warmup, constant schedule, no gradient clipping
									_a     :      List[Any]        =							AdamW(params=[w]	,lr=2E-1	,weight_decay=0.0		)
									for _ in range(100		):
											_a     :      int        =							criterion(__lowerCAmelCase	,__lowerCAmelCase		)
											loss.backward()
											optimizer.step()
											w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
											w.grad.zero_()
									self.assertListAlmostEqual(w.tolist()	,[0.4, 0.2, -0.5]	,tol=1E-2		)







							def 							__lowercase			(							self					:       Union[str, Any]		):
									_a     :      List[str]        =							torch.tensor([0.1, -0.2, -0.1]	,requires_grad=__lowerCAmelCase		)
									_a     :      Any        =							torch.tensor([0.4, 0.2, -0.5]		)
									_a     :      Optional[Any]        =							nn.MSELoss()
									# No warmup, constant schedule, no gradient clipping
									_a     :      Dict        =							Adafactor(
									    params=[w]	,lr=1E-2	,eps=(1E-30, 1E-3)	,clip_threshold=1.0	,decay_rate=-0.8	,betaa=__lowerCAmelCase	,weight_decay=0.0	,relative_step=__lowerCAmelCase	,scale_parameter=__lowerCAmelCase	,warmup_init=__lowerCAmelCase	,)
									for _ in range(1000		):
											_a     :      str        =							criterion(__lowerCAmelCase	,__lowerCAmelCase		)
											loss.backward()
											optimizer.step()
											w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
											w.grad.zero_()
									self.assertListAlmostEqual(w.tolist()	,[0.4, 0.2, -0.5]	,tol=1E-2		)

@require_torch
class 	__magic_name__						(    unittest.TestCase   ):
							lowerCAmelCase       :    Union[str, Any]   				=					nn.Linear(5_0			,						5_0   ) if is_torch_available() else None
							lowerCAmelCase       :    Union[str, Any]   				=					AdamW(m.parameters()			,						lr=10.0   ) if is_torch_available() else None
							lowerCAmelCase       :    Dict   				=					1_0







							def 							__lowercase			(							self					:       Any	,_UpperCAmelCase					:       Dict	,_UpperCAmelCase					:       List[Any]	,_UpperCAmelCase					:       Tuple	,_UpperCAmelCase					:       Optional[Any]=None		):
									self.assertEqual(len(__lowerCAmelCase		)	,len(__lowerCAmelCase		)		)
									for a, b in zip(__lowerCAmelCase	,__lowerCAmelCase		):
											self.assertAlmostEqual(__lowerCAmelCase	,__lowerCAmelCase	,delta=__lowerCAmelCase	,msg=__lowerCAmelCase		)







							def 							__lowercase			(							self					:       Any		):
									_a     :      Optional[int]        =							{'num_warmup_steps': 2, 'num_training_steps': 10}
									# schedulers doct format
									# function: (sched_args_dict, expected_learning_rates)
									_a     :      Optional[int]        =							{
									    get_constant_schedule: ({}, [10.0] * self.num_steps),
									    get_constant_schedule_with_warmup: (
									        {'num_warmup_steps': 4},
									        [0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0],
									    ),
									    get_linear_schedule_with_warmup: (
									        {**common_kwargs},
									        [0.0, 5.0, 10.0, 8.75, 7.5, 6.25, 5.0, 3.75, 2.5, 1.25],
									    ),
									    get_cosine_schedule_with_warmup: (
									        {**common_kwargs},
									        [0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38],
									    ),
									    get_cosine_with_hard_restarts_schedule_with_warmup: (
									        {**common_kwargs, 'num_cycles': 2},
									        [0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46],
									    ),
									    get_polynomial_decay_schedule_with_warmup: (
									        {**common_kwargs, 'power': 2.0, 'lr_end': 1E-7},
									        [0.0, 5.0, 10.0, 7.6_56, 5.6_25, 3.9_06, 2.5, 1.4_06, 0.6_25, 0.1_56],
									    ),
									    get_inverse_sqrt_schedule: (
									        {'num_warmup_steps': 2},
									        [0.0, 5.0, 10.0, 8.1_65, 7.0_71, 6.3_25, 5.7_74, 5.3_45, 5.0, 4.7_14],
									    ),
									}

									for scheduler_func, data in scheds.items():
											_a       ,					_a     :      List[str]        =							data

											_a     :      Dict        =							scheduler_func(self.optimizer	,**__lowerCAmelCase		)
											self.assertEqual(len([scheduler.get_lr()[0]]		)	,1		)
											_a     :      Dict        =							unwrap_schedule(__lowerCAmelCase	,self.num_steps		)
											self.assertListAlmostEqual(
											    __lowerCAmelCase	,__lowerCAmelCase	,tol=1E-2	,msg=F"""failed for {scheduler_func} in normal scheduler"""	,)

											_a     :      List[Any]        =							scheduler_func(self.optimizer	,**__lowerCAmelCase		)
											if scheduler_func.__name__ != "get_constant_schedule":
													LambdaScheduleWrapper.wrap_scheduler(__lowerCAmelCase		)  # wrap to test picklability of the schedule
											_a     :      Optional[Any]        =							unwrap_and_save_reload_schedule(__lowerCAmelCase	,self.num_steps		)
											self.assertListEqual(__lowerCAmelCase	,__lowerCAmelCase	,msg=F"""failed for {scheduler_func} in save and reload"""		)

class 	__magic_name__						:







							def __init__(							self					:       Tuple	,_UpperCAmelCase					:       Optional[int]		):
									_a     :      Tuple        =							fn







							def __call__(							self					:       Tuple	,*_UpperCAmelCase					:       int	,**_UpperCAmelCase					:       int		):
									return self.fn(*__lowerCAmelCase	,**__lowerCAmelCase		)







							@classmethod
							def 							__lowercase			(							self					:       List[Any]	,_UpperCAmelCase					:       Union[str, Any]		):
									_a     :      Optional[Any]        =							list(map(self	,scheduler.lr_lambdas		)		)





"""simple docstring"""






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


__lowerCAmelCase    :       List[str]					     =logging.get_logger(__name__)

__lowerCAmelCase    :       str					     ={
    """facebook/dpr-ctx_encoder-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-question_encoder-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-reader-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-ctx_encoder-multiset-base""": (
        """https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/config.json"""
    ),
    """facebook/dpr-question_encoder-multiset-base""": (
        """https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/config.json"""
    ),
    """facebook/dpr-reader-multiset-base""": (
        """https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/config.json"""
    ),
}





class  _A						(       lowerCAmelCase	):
      snake_case__			:   Tuple							  =	'dpr'

      def __init__(						self		,       __lowerCAmelCase=3_0522		,       __lowerCAmelCase=768		,       __lowerCAmelCase=12		,       __lowerCAmelCase=12		,       __lowerCAmelCase=3072		,       __lowerCAmelCase="gelu"		,       __lowerCAmelCase=0.1		,       __lowerCAmelCase=0.1		,       __lowerCAmelCase=512		,       __lowerCAmelCase=2		,       __lowerCAmelCase=0.0_2		,       __lowerCAmelCase=1E-12		,       __lowerCAmelCase=0		,       __lowerCAmelCase="absolute"		,       __lowerCAmelCase = 0		,       **__lowerCAmelCase		,       ):







            """simple docstring"""



            super().__init__(pad_token_id=__lowerCAmelCase		,       **__lowerCAmelCase						)

            lowercase														=	vocab_size
            lowercase														=	hidden_size
            lowercase														=	num_hidden_layers
            lowercase														=	num_attention_heads
            lowercase														=	hidden_act
            lowercase														=	intermediate_size
            lowercase														=	hidden_dropout_prob
            lowercase														=	attention_probs_dropout_prob
            lowercase														=	max_position_embeddings
            lowercase														=	type_vocab_size
            lowercase														=	initializer_range
            lowercase														=	layer_norm_eps
            lowercase														=	projection_dim
            lowercase														=	position_embedding_type



