

'''simple docstring'''

import inspect
import warnings

import torch

from .state import AcceleratorState, GradientState
from .utils import DistributedType, honor_type, is_tpu_available


if is_tpu_available(check_device=False):
       import torch_xla.core.xla_model as xm

def 				__UpperCamelCase							(	lowercase__						:  int,					lowercase__						:  List[Any] ):






     '''simple docstring'''





     if isinstance(lowercase__,					(list, tuple) ):
          return honor_type(lowercase__,					(move_to_device(lowercase__,					lowercase__ ) for t in state) )
     elif isinstance(lowercase__,					lowercase__ ):
          return type(lowercase__ )({k: move_to_device(lowercase__,					lowercase__ ) for k, v in state.items()} )
     elif isinstance(lowercase__,					torch.Tensor ):
          return state.to(lowercase__ )
     return state







class   lowerCAmelCase     (      torch.optim.Optimizer   ):







       def __init__(							self			: int						,      __lowercase			: Dict						,      __lowercase			: Tuple=True						,      __lowercase			: List[str]=None      ):







            """simple docstring"""







            __lowercase					=optimizer
            __lowercase					=scaler
            __lowercase					=AcceleratorState()
            __lowercase					=GradientState()
            __lowercase					=device_placement
            __lowercase					=False
            __lowercase					=None

            # Handle device placement
            if device_placement:
                 __lowercase					=self.optimizer.state_dict()
                 if self.accelerator_state.distributed_type == DistributedType.TPU:
                      xm.send_cpu_data_to_device(__lowercase						,      self.accelerator_state.device      )
                 else:
                      __lowercase					=move_to_device(__lowercase						,      self.accelerator_state.device      )
                 self.optimizer.load_state_dict(__lowercase      )







       @property
       def 			snake_case   (							self			: int      ):







            """simple docstring"""







            return self.optimizer.state







       @state.setter
       def 			snake_case   (							self			: Union[str, Any]						,      __lowercase			: Tuple      ):







            """simple docstring"""







            __lowercase					=state







       @property
       def 			snake_case   (							self			: Optional[int]      ):







            """simple docstring"""







            return self.optimizer.param_groups







       @param_groups.setter
       def 			snake_case   (							self			: List[str]						,      __lowercase			: Tuple      ):







            """simple docstring"""







            __lowercase					=param_groups







       @property
       def 			snake_case   (							self			: List[Any]      ):







            """simple docstring"""







            return self.optimizer.defaults







       @defaults.setter
       def 			snake_case   (							self			: Tuple						,      __lowercase			: List[Any]      ):







            """simple docstring"""







            __lowercase					=defaults







       def 			snake_case   (							self			: Tuple						,      __lowercase			: str      ):







            """simple docstring"""







            self.optimizer.add_param_group(__lowercase      )







       def 			snake_case   (							self			: Optional[Any]						,      __lowercase			: Optional[Any]      ):







            """simple docstring"""







            if self.accelerator_state.distributed_type == DistributedType.TPU and self.device_placement:
                 xm.send_cpu_data_to_device(__lowercase						,      self.accelerator_state.device      )
            self.optimizer.load_state_dict(__lowercase      )







       def 			snake_case   (							self			: Dict      ):







            """simple docstring"""







            return self.optimizer.state_dict()







       def 			snake_case   (							self			: List[str]						,      __lowercase			: Tuple=None      ):







            """simple docstring"""







            if self.gradient_state.sync_gradients:
                 __lowercase					='set_to_none' in inspect.signature(self.optimizer.zero_grad      ).parameters
                 if accept_arg:
                      if set_to_none is None:
                           __lowercase					=False
                      self.optimizer.zero_grad(set_to_none=__lowercase      )
                 else:
                      if set_to_none is not None:
                           raise ValueError('`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer.'      )
                      self.optimizer.zero_grad()







       def 			snake_case   (							self			: Optional[int]						,      __lowercase			: int=None      ):







            """simple docstring"""







            if self.gradient_state.sync_gradients:
                 if self.accelerator_state.distributed_type == DistributedType.TPU:
                      __lowercase					={'closure': closure} if closure is not None else {}
                      xm.optimizer_step(self.optimizer						,      optimizer_args=__lowercase      )
                 elif self.scaler is not None:
                      __lowercase					=False
                      if self._last_scale is None:
                           # `get_scale` is an async operation requiring full synchronization
                           # on CPU and GPUs before finishing. As a result, we store away
                           # the prior one to reduce the call overhead
                           __lowercase					=self.scaler.get_scale()
                           __lowercase					=True
                      self.scaler.step(self.optimizer						,      __lowercase      )
                      self.scaler.update()
                      __lowercase					=self.scaler.get_scale()
                      if not new_scale:
                           # If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
                           __lowercase					=scale_after < self._last_scale
                      __lowercase					=scale_after

                 else:
                      self.optimizer.step(__lowercase      )







       def 			snake_case   (							self			: Union[str, Any]						,      __lowercase			: List[str]      ):







            """simple docstring"""







            for param_group in self.optimizer.param_groups:
                 __lowercase					=[parameters_map.get(__lowercase						,      __lowercase      ) for p in param_group['params']]







       @property
       def 			snake_case   (							self			: Dict      ):







            """simple docstring"""







            warnings.warn(
                'The `is_overflow` property is deprecated and will be removed in version 1.0 of Accelerate use '
                '`optimizer.step_was_skipped` instead.'						,      __lowercase						,      )
            return self._is_overflow







       @property
       def 			snake_case   (							self			: Dict      ):







            """simple docstring"""







            return self._is_overflow







       def __getstate__(							self			: Optional[int]      ):







            """simple docstring"""







            return self.__dict__.copy()







       def __setstate__(							self			: Any						,      __lowercase			: List[str]      ):







            """simple docstring"""







            self.__dict__.update(__lowercase      )





'''simple docstring'''

import tempfile
import unittest

from transformers import DPRConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import DPRContextEncoder, DPRQuestionEncoder, DPRReader, DPRReaderTokenizer
       from transformers.models.dpr.modeling_dpr import (
           DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
           DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
           DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
       )


class   lowerCAmelCase     :







       def __init__(							self			: List[str]						,      __lowercase			: List[Any]						,      __lowercase			: Union[str, Any]=13						,      __lowercase			: Optional[int]=7						,      __lowercase			: Union[str, Any]=False						,      __lowercase			: str=True						,      __lowercase			: Dict=True						,      __lowercase			: Optional[Any]=True						,      __lowercase			: int=99						,      __lowercase			: Optional[Any]=32						,      __lowercase			: Optional[int]=5						,      __lowercase			: List[str]=4						,      __lowercase			: Any=37						,      __lowercase			: str="gelu"						,      __lowercase			: str=0.1						,      __lowercase			: str=0.1						,      __lowercase			: str=512						,      __lowercase			: Dict=16						,      __lowercase			: Dict=2						,      __lowercase			: List[str]=0.0_2						,      __lowercase			: Union[str, Any]=3						,      __lowercase			: List[Any]=4						,      __lowercase			: List[str]=None						,      __lowercase			: List[Any]=0						,      ):







            """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 			snake_case   (							self			: Any      ):







            """simple docstring"""







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

            __lowercase					=None
            if self.use_input_mask:
                 __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					=self.get_config()

            return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







       def 			snake_case   (							self			: Dict      ):







            """simple docstring"""







            return DPRConfig(
                projection_dim=self.projection_dim						,      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						,      initializer_range=self.initializer_range						,      )







       def 			snake_case   (							self			: Union[str, Any]						,      __lowercase			: int						,      __lowercase			: Tuple						,      __lowercase			: Any						,      __lowercase			: Union[str, Any]						,      __lowercase			: List[Any]						,      __lowercase			: Tuple						,      __lowercase			: List[str]      ):







            """simple docstring"""







            __lowercase					=DPRContextEncoder(config=__lowercase      )
            model.to(__lowercase      )
            model.eval()
            __lowercase					=model(__lowercase						,      attention_mask=__lowercase						,      token_type_ids=__lowercase      )
            __lowercase					=model(__lowercase						,      token_type_ids=__lowercase      )
            __lowercase					=model(__lowercase      )
            self.parent.assertEqual(result.pooler_output.shape						,      (self.batch_size, self.projection_dim or self.hidden_size)      )







       def 			snake_case   (							self			: Union[str, Any]						,      __lowercase			: Optional[int]						,      __lowercase			: Optional[int]						,      __lowercase			: Union[str, Any]						,      __lowercase			: Any						,      __lowercase			: Optional[Any]						,      __lowercase			: Tuple						,      __lowercase			: Dict      ):







            """simple docstring"""







            __lowercase					=DPRQuestionEncoder(config=__lowercase      )
            model.to(__lowercase      )
            model.eval()
            __lowercase					=model(__lowercase						,      attention_mask=__lowercase						,      token_type_ids=__lowercase      )
            __lowercase					=model(__lowercase						,      token_type_ids=__lowercase      )
            __lowercase					=model(__lowercase      )
            self.parent.assertEqual(result.pooler_output.shape						,      (self.batch_size, self.projection_dim or self.hidden_size)      )







       def 			snake_case   (							self			: List[str]						,      __lowercase			: Optional[Any]						,      __lowercase			: Tuple						,      __lowercase			: str						,      __lowercase			: Tuple						,      __lowercase			: str						,      __lowercase			: List[Any]						,      __lowercase			: Optional[Any]      ):







            """simple docstring"""







            __lowercase					=DPRReader(config=__lowercase      )
            model.to(__lowercase      )
            model.eval()
            __lowercase					=model(
                __lowercase						,      attention_mask=__lowercase						,      )

            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 			snake_case   (							self			: List[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_torch
class   lowerCAmelCase     (      A   ,			A   ,			unittest.TestCase   ):
       lowerCAmelCase_		= (
           (
               DPRContextEncoder,
               DPRQuestionEncoder,
               DPRReader,
           )
           if is_torch_available()
           else ()
       )
       lowerCAmelCase_		= {"feature-extraction": DPRQuestionEncoder} if is_torch_available() else {}

       lowerCAmelCase_		= False
       lowerCAmelCase_		= False  # why?
       lowerCAmelCase_		= False
       lowerCAmelCase_		= False







       def 			snake_case   (							self			: int      ):







            """simple docstring"""







            __lowercase					=DPRModelTester(self      )
            __lowercase					=ConfigTester(self						,      config_class=__lowercase						,      hidden_size=37      )







       def 			snake_case   (							self			: Any      ):







            """simple docstring"""







            self.config_tester.run_common_tests()







       def 			snake_case   (							self			: List[str]      ):







            """simple docstring"""







            __lowercase					=self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_context_encoder(*__lowercase      )







       def 			snake_case   (							self			: Optional[Any]      ):







            """simple docstring"""







            __lowercase					=self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_question_encoder(*__lowercase      )







       def 			snake_case   (							self			: Any      ):







            """simple docstring"""







            __lowercase					=self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_reader(*__lowercase      )







       def 			snake_case   (							self			: Optional[int]      ):







            """simple docstring"""







            __lowercase					=self.model_tester.prepare_config_and_inputs()[0]

            __lowercase					=DPRQuestionEncoder(config=__lowercase      )
            model.to(__lowercase      )
            model.eval()

            with tempfile.TemporaryDirectory() as tmp_dirname:
                 model.save_pretrained(__lowercase      )
                 __lowercase					=DPRQuestionEncoder.from_pretrained(__lowercase						,      projection_dim=512      )

            self.assertIsNotNone(__lowercase      )







       @slow
       def 			snake_case   (							self			: int      ):







            """simple docstring"""







            for model_name in DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 __lowercase					=DPRContextEncoder.from_pretrained(__lowercase      )
                 self.assertIsNotNone(__lowercase      )

            for model_name in DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 __lowercase					=DPRContextEncoder.from_pretrained(__lowercase      )
                 self.assertIsNotNone(__lowercase      )

            for model_name in DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 __lowercase					=DPRQuestionEncoder.from_pretrained(__lowercase      )
                 self.assertIsNotNone(__lowercase      )

            for model_name in DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 __lowercase					=DPRReader.from_pretrained(__lowercase      )
                 self.assertIsNotNone(__lowercase      )






@require_torch
class   lowerCAmelCase     (      unittest.TestCase   ):







       @slow
       def 			snake_case   (							self			: int      ):







            """simple docstring"""







            __lowercase					=DPRQuestionEncoder.from_pretrained('facebook/dpr-question_encoder-single-nq-base'						,      return_dict=__lowercase      )
            model.to(__lowercase      )

            __lowercase					=torch.tensor(
                [[101, 7592, 1010, 2003, 2026, 3899, 10140, 1029, 102]]						,      dtype=torch.long						,      device=__lowercase      )  # [CLS] hello, is my dog cute? [SEP]
            __lowercase					=model(__lowercase      )[0]  # embedding shape = (1, 768)
            # compare the actual values for a slice.
            __lowercase					=torch.tensor(
                [
                    [
                        0.0_3_2_3_6_2_5_3,
                        0.1_2_7_5_3_3_3_5,
                        0.1_6_8_1_8_5_0_9,
                        0.0_0_2_7_9_7_8_6,
                        0.3_8_9_6_9_3_3,
                        0.2_4_2_6_4_9_4_5,
                        0.2_1_7_8_9_7_1,
                        -0.0_2_3_3_5_2_2_7,
                        -0.0_8_4_8_1_9_5_9,
                        -0.1_4_3_2_4_1_1_7,
                    ]
                ]						,      dtype=torch.float						,      device=__lowercase						,      )
            self.assertTrue(torch.allclose(output[:, :10]						,      __lowercase						,      atol=1E-4      )      )







       @slow
       def 			snake_case   (							self			: Union[str, Any]      ):







            """simple docstring"""







            __lowercase					=DPRReaderTokenizer.from_pretrained('facebook/dpr-reader-single-nq-base'      )
            __lowercase					=DPRReader.from_pretrained('facebook/dpr-reader-single-nq-base'      )
            model.to(__lowercase      )

            __lowercase					=tokenizer(
                questions='What is love ?'						,      titles='Haddaway'						,      texts='What Is Love is a song recorded by the artist Haddaway'						,      padding=__lowercase						,      return_tensors='pt'						,      )
            encoded_inputs.to(__lowercase      )

            __lowercase					=model(**__lowercase      )

            # compare the actual values for a slice.
            __lowercase					=torch.tensor(
                [[-1_0.3_0_0_5, -1_0.7_7_6_5, -1_1.4_8_7_2, -1_1.6_8_4_1, -1_1.9_3_1_2, -1_0.3_0_0_2, -9.8_5_4_4, -1_1.7_3_7_8, -1_2.0_8_2_1, -1_0.2_9_7_5]]						,      dtype=torch.float						,      device=__lowercase						,      )

            __lowercase					=torch.tensor(
                [[-1_1.0_6_8_4, -1_1.7_0_4_1, -1_1.5_3_9_7, -1_0.3_4_6_5, -1_0.8_7_9_1, -6.8_4_4_3, -1_1.9_9_5_9, -1_1.0_3_6_4, -1_0.0_0_9_6, -6.8_4_0_5]]						,      dtype=torch.float						,      device=__lowercase						,      )
            self.assertTrue(torch.allclose(outputs.start_logits[:, :10]						,      __lowercase						,      atol=1E-4      )      )
            self.assertTrue(torch.allclose(outputs.end_logits[:, :10]						,      __lowercase						,      atol=1E-4      )      )


