


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       :






  """simple docstring"""





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




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



  def    lowerCAmelCase__					(       self    :    Tuple					)				->	Tuple:




     """simple docstring"""
     snake_case_         =   ids_tensor([self.batch_size, self.seq_length]      ,			self.vocab_size					)

     snake_case_         =   None
     if self.use_input_mask:
        snake_case_         =   random_attention_mask([self.batch_size, self.seq_length]					)

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

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

     snake_case_         =   self.get_config()

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



  def    lowerCAmelCase__					(       self    :    Dict					)				->	List[str]:




     """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    lowerCAmelCase__					(       self    :    List[Any]      ,			_lowerCAmelCase    :    Optional[Any]      ,			_lowerCAmelCase    :    List[Any]      ,			_lowerCAmelCase    :    Dict      ,			_lowerCAmelCase    :    str      ,			_lowerCAmelCase    :    int      ,			_lowerCAmelCase    :    List[Any]      ,			_lowerCAmelCase    :    Dict					)				->	Tuple:




     """simple docstring"""
     snake_case_         =   DPRContextEncoder(config=_lowerCAmelCase					)
     model.to(_lowerCAmelCase					)
     model.eval()
     snake_case_         =   model(_lowerCAmelCase      ,			attention_mask=_lowerCAmelCase      ,			token_type_ids=_lowerCAmelCase					)
     snake_case_         =   model(_lowerCAmelCase      ,			token_type_ids=_lowerCAmelCase					)
     snake_case_         =   model(_lowerCAmelCase					)
     self.parent.assertEqual(result.pooler_output.shape      ,			(self.batch_size, self.projection_dim or self.hidden_size)					)



  def    lowerCAmelCase__					(       self    :    Optional[Any]      ,			_lowerCAmelCase    :    Tuple      ,			_lowerCAmelCase    :    Dict      ,			_lowerCAmelCase    :    Union[str, Any]      ,			_lowerCAmelCase    :    int      ,			_lowerCAmelCase    :    Dict      ,			_lowerCAmelCase    :    Dict      ,			_lowerCAmelCase    :    int					)				->	int:




     """simple docstring"""
     snake_case_         =   DPRQuestionEncoder(config=_lowerCAmelCase					)
     model.to(_lowerCAmelCase					)
     model.eval()
     snake_case_         =   model(_lowerCAmelCase      ,			attention_mask=_lowerCAmelCase      ,			token_type_ids=_lowerCAmelCase					)
     snake_case_         =   model(_lowerCAmelCase      ,			token_type_ids=_lowerCAmelCase					)
     snake_case_         =   model(_lowerCAmelCase					)
     self.parent.assertEqual(result.pooler_output.shape      ,			(self.batch_size, self.projection_dim or self.hidden_size)					)



  def    lowerCAmelCase__					(       self    :    str      ,			_lowerCAmelCase    :    Optional[int]      ,			_lowerCAmelCase    :    Any      ,			_lowerCAmelCase    :    List[str]      ,			_lowerCAmelCase    :    List[str]      ,			_lowerCAmelCase    :    Any      ,			_lowerCAmelCase    :    Optional[int]      ,			_lowerCAmelCase    :    List[str]					)				->	str:




     """simple docstring"""
     snake_case_         =   DPRReader(config=_lowerCAmelCase					)
     model.to(_lowerCAmelCase					)
     model.eval()
     snake_case_         =   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    lowerCAmelCase__					(       self    :    int					)				->	str:




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






@require_torch
class       __lowerCAmelCase       (	a	,					a	,					unittest.TestCase   ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE						      =    (
      (
          DPRContextEncoder,
          DPRQuestionEncoder,
          DPRReader,
      )
      if is_torch_available()
      else ()
  )
  _SCREAMING_SNAKE_CASE						      =    {'feature-extraction': DPRQuestionEncoder} if is_torch_available() else {}

  _SCREAMING_SNAKE_CASE						      =    False
  _SCREAMING_SNAKE_CASE						      =    False  # why?
  _SCREAMING_SNAKE_CASE						      =    False
  _SCREAMING_SNAKE_CASE						      =    False



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




     """simple docstring"""
     snake_case_         =   DPRModelTester(self					)
     snake_case_         =   ConfigTester(self      ,			config_class=_lowerCAmelCase      ,			hidden_size=3_7					)



  def    lowerCAmelCase__					(       self    :    Any					)				->	int:




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



  def    lowerCAmelCase__					(       self    :    str					)				->	Tuple:




     """simple docstring"""
     snake_case_         =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_context_encoder(*_lowerCAmelCase					)



  def    lowerCAmelCase__					(       self    :    Any					)				->	int:




     """simple docstring"""
     snake_case_         =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_question_encoder(*_lowerCAmelCase					)



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




     """simple docstring"""
     snake_case_         =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_reader(*_lowerCAmelCase					)



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




     """simple docstring"""
     snake_case_         =   self.model_tester.prepare_config_and_inputs()[0]

     snake_case_         =   DPRQuestionEncoder(config=_lowerCAmelCase					)
     model.to(_lowerCAmelCase					)
     model.eval()

     with tempfile.TemporaryDirectory() as tmp_dirname:
        model.save_pretrained(_lowerCAmelCase					)
        snake_case_         =   DPRQuestionEncoder.from_pretrained(_lowerCAmelCase      ,			projection_dim=5_1_2					)

     self.assertIsNotNone(_lowerCAmelCase					)



  @slow
  def    lowerCAmelCase__					(       self    :    List[Any]					)				->	Optional[int]:




     """simple docstring"""
     for model_name in DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        snake_case_         =   DPRContextEncoder.from_pretrained(_lowerCAmelCase					)
        self.assertIsNotNone(_lowerCAmelCase					)

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

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

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






@require_torch
class       __lowerCAmelCase       (	unittest.TestCase   ):






  """simple docstring"""





  @slow
  def    lowerCAmelCase__					(       self    :    Union[str, Any]					)				->	Optional[Any]:




     """simple docstring"""
     snake_case_         =   DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base"      ,			return_dict=_lowerCAmelCase					)
     model.to(_lowerCAmelCase					)

     snake_case_         =   torch.tensor(
         [[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]]      ,			dtype=torch.long      ,			device=_lowerCAmelCase					)  # [CLS] hello, is my dog cute? [SEP]
     snake_case_         =   model(_lowerCAmelCase					)[0]  # embedding shape = (1, 768)
     # compare the actual values for a slice.
     snake_case_         =   torch.tensor(
         [
             [
                 0.03_236_253,
                 0.12_753_335,
                 0.16_818_509,
                 0.00_279_786,
                 0.3_896_933,
                 0.24_264_945,
                 0.2_178_971,
                 -0.02_335_227,
                 -0.08_481_959,
                 -0.14_324_117,
             ]
         ]      ,			dtype=torch.float      ,			device=_lowerCAmelCase      ,			)
     self.assertTrue(torch.allclose(output[:, :1_0]      ,			_lowerCAmelCase      ,			atol=1e-4					)					)



  @slow
  def    lowerCAmelCase__					(       self    :    List[str]					)				->	Optional[int]:




     """simple docstring"""
     snake_case_         =   DPRReaderTokenizer.from_pretrained("facebook/dpr-reader-single-nq-base"					)
     snake_case_         =   DPRReader.from_pretrained("facebook/dpr-reader-single-nq-base"					)
     model.to(_lowerCAmelCase					)

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

     snake_case_         =   model(**_lowerCAmelCase					)

     # compare the actual values for a slice.
     snake_case_         =   torch.tensor(
         [[-10.3_005, -10.7_765, -11.4_872, -11.6_841, -11.9_312, -10.3_002, -9.8_544, -11.7_378, -12.0_821, -10.2_975]]      ,			dtype=torch.float      ,			device=_lowerCAmelCase      ,			)

     snake_case_         =   torch.tensor(
         [[-11.0_684, -11.7_041, -11.5_397, -10.3_465, -10.8_791, -6.8_443, -11.9_959, -11.0_364, -10.0_096, -6.8_405]]      ,			dtype=torch.float      ,			device=_lowerCAmelCase      ,			)
     self.assertTrue(torch.allclose(outputs.start_logits[:, :1_0]      ,			_lowerCAmelCase      ,			atol=1e-4					)					)
     self.assertTrue(torch.allclose(outputs.end_logits[:, :1_0]      ,			_lowerCAmelCase      ,			atol=1e-4					)					)








from __future__ import annotations



def      _lowerCAmelCase  (							lowerCAmelCase_							:int							)->list[int]:
   '''simple docstring'''






   assert (
       isinstance(lowerCAmelCase_						,  lowerCAmelCase_							) and mask > 0
   ), F'''mask needs to be positive integer, your input {mask}'''
   snake_case_         =   []
   snake_case_         =   mask

   while submask:
      all_submasks.append(lowerCAmelCase_							)
      snake_case_         =   (submask - 1) & mask

   return all_submasks


if __name__ == "__main__":
   import doctest

   doctest.testmod()




