


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						,	UpperCAmelCase						,	UpperCAmelCase=1_3						,	UpperCAmelCase=7						,	UpperCAmelCase=False						,	UpperCAmelCase=True						,	UpperCAmelCase=True						,	UpperCAmelCase=True						,	UpperCAmelCase=9_9						,	UpperCAmelCase=3_2						,	UpperCAmelCase=5						,	UpperCAmelCase=4						,	UpperCAmelCase=3_7						,	UpperCAmelCase="gelu"						,	UpperCAmelCase=0.1						,	UpperCAmelCase=0.1						,	UpperCAmelCase=5_1_2						,	UpperCAmelCase=1_6						,	UpperCAmelCase=2						,	UpperCAmelCase=0.02						,	UpperCAmelCase=3						,	UpperCAmelCase=4						,	UpperCAmelCase=None						,	UpperCAmelCase=0						,	)	->  str:
      _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 	__A   (self			)	->  Tuple:
      _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 	__A   (self			)	->  Dict:
      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 	__A   (self						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase			)	->  List[Any]:
      _lowercase				 =DPRContextEncoder(config=UpperCAmelCase			)
      model.to(UpperCAmelCase			)
      model.eval()
      _lowercase				 =model(UpperCAmelCase						,	attention_mask=UpperCAmelCase						,	token_type_ids=UpperCAmelCase			)
      _lowercase				 =model(UpperCAmelCase						,	token_type_ids=UpperCAmelCase			)
      _lowercase				 =model(UpperCAmelCase			)
      self.parent.assertEqual(result.pooler_output.shape						,	(self.batch_size, self.projection_dim or self.hidden_size)			)





     def 	__A   (self						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase			)	->  Dict:
      _lowercase				 =DPRQuestionEncoder(config=UpperCAmelCase			)
      model.to(UpperCAmelCase			)
      model.eval()
      _lowercase				 =model(UpperCAmelCase						,	attention_mask=UpperCAmelCase						,	token_type_ids=UpperCAmelCase			)
      _lowercase				 =model(UpperCAmelCase						,	token_type_ids=UpperCAmelCase			)
      _lowercase				 =model(UpperCAmelCase			)
      self.parent.assertEqual(result.pooler_output.shape						,	(self.batch_size, self.projection_dim or self.hidden_size)			)





     def 	__A   (self						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase						,	UpperCAmelCase			)	->  str:
      _lowercase				 =DPRReader(config=UpperCAmelCase			)
      model.to(UpperCAmelCase			)
      model.eval()
      _lowercase				 =model(
          UpperCAmelCase						,	attention_mask=UpperCAmelCase						,	)

      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 	__A   (self			)	->  int:
      _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__  (			lowerCAmelCase  ,				lowerCAmelCase  ,				unittest.TestCase):
     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 	__A   (self			)	->  int:
      _lowercase				 =DPRModelTester(self			)
      _lowercase				 =ConfigTester(self						,	config_class=UpperCAmelCase						,	hidden_size=3_7			)





     def 	__A   (self			)	->  Tuple:
      self.config_tester.run_common_tests()





     def 	__A   (self			)	->  List[Any]:
      _lowercase				 =self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_context_encoder(*UpperCAmelCase			)





     def 	__A   (self			)	->  Dict:
      _lowercase				 =self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_question_encoder(*UpperCAmelCase			)





     def 	__A   (self			)	->  Optional[int]:
      _lowercase				 =self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_reader(*UpperCAmelCase			)





     def 	__A   (self			)	->  Optional[int]:
      _lowercase				 =self.model_tester.prepare_config_and_inputs()[0]

      _lowercase				 =DPRQuestionEncoder(config=UpperCAmelCase			)
      model.to(UpperCAmelCase			)
      model.eval()

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

      self.assertIsNotNone(UpperCAmelCase			)





     @slow
     def 	__A   (self			)	->  Dict:
      for model_name in DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       _lowercase				 =DPRContextEncoder.from_pretrained(UpperCAmelCase			)
       self.assertIsNotNone(UpperCAmelCase			)

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

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

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







@require_torch
class  lowerCamelCase__  (			unittest.TestCase):





     @slow
     def 	__A   (self			)	->  List[str]:
      _lowercase				 =DPRQuestionEncoder.from_pretrained('''facebook/dpr-question_encoder-single-nq-base'''						,	return_dict=UpperCAmelCase			)
      model.to(UpperCAmelCase			)

      _lowercase				 =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=UpperCAmelCase			)  # [CLS] hello, is my dog cute? [SEP]
      _lowercase				 =model(UpperCAmelCase			)[0]  # embedding shape = (1, 768)
      # compare the actual values for a slice.
      _lowercase				 =torch.tensor(
          [
              [
                  0.0323_6253,
                  0.1275_3335,
                  0.1681_8509,
                  0.0027_9786,
                  0.389_6933,
                  0.2426_4945,
                  0.217_8971,
                  -0.0233_5227,
                  -0.0848_1959,
                  -0.1432_4117,
              ]
          ]						,	dtype=torch.float						,	device=UpperCAmelCase						,	)
      self.assertTrue(torch.allclose(output[:, :1_0]						,	UpperCAmelCase						,	atol=1e-4			)			)





     @slow
     def 	__A   (self			)	->  Dict:
      _lowercase				 =DPRReaderTokenizer.from_pretrained('''facebook/dpr-reader-single-nq-base'''			)
      _lowercase				 =DPRReader.from_pretrained('''facebook/dpr-reader-single-nq-base'''			)
      model.to(UpperCAmelCase			)

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

      _lowercase				 =model(**UpperCAmelCase			)

      # compare the actual values for a slice.
      _lowercase				 =torch.tensor(
          [[-10.3005, -10.7765, -11.4872, -11.6841, -11.9312, -10.3002, -9.8544, -11.7378, -12.0821, -10.2975]]						,	dtype=torch.float						,	device=UpperCAmelCase						,	)

      _lowercase				 =torch.tensor(
          [[-11.0684, -11.7041, -11.5397, -10.3465, -10.8791, -6.8443, -11.9959, -11.0364, -10.0096, -6.8405]]						,	dtype=torch.float						,	device=UpperCAmelCase						,	)
      self.assertTrue(torch.allclose(outputs.start_logits[:, :1_0]						,	UpperCAmelCase						,	atol=1e-4			)			)
      self.assertTrue(torch.allclose(outputs.end_logits[:, :1_0]						,	UpperCAmelCase						,	atol=1e-4			)			)





from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCAmelCase__        =				logging.get_logger(__name__)

UpperCAmelCase__        =				{
    '''google/rembert''': '''https://huggingface.co/google/rembert/resolve/main/config.json''',
    # See all RemBERT models at https://huggingface.co/models?filter=rembert
}







class  lowerCamelCase__  (			lowerCAmelCase):
     SCREAMING_SNAKE_CASE__      =     '''rembert'''





     def __init__(self						,	UpperCAmelCase=2_5_0_3_0_0						,	UpperCAmelCase=1_1_5_2						,	UpperCAmelCase=3_2						,	UpperCAmelCase=1_8						,	UpperCAmelCase=2_5_6						,	UpperCAmelCase=1_6_6_4						,	UpperCAmelCase=4_6_0_8						,	UpperCAmelCase="gelu"						,	UpperCAmelCase=0.0						,	UpperCAmelCase=0.0						,	UpperCAmelCase=0.1						,	UpperCAmelCase=5_1_2						,	UpperCAmelCase=2						,	UpperCAmelCase=0.02						,	UpperCAmelCase=1e-12						,	UpperCAmelCase=True						,	UpperCAmelCase=0						,	UpperCAmelCase=3_1_2						,	UpperCAmelCase=3_1_3						,	**UpperCAmelCase						,	)	->  Tuple:
      super().__init__(pad_token_id=UpperCAmelCase						,	bos_token_id=UpperCAmelCase						,	eos_token_id=UpperCAmelCase						,	**UpperCAmelCase			)

      _lowercase				 =vocab_size
      _lowercase				 =input_embedding_size
      _lowercase				 =output_embedding_size
      _lowercase				 =max_position_embeddings
      _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				 =classifier_dropout_prob
      _lowercase				 =initializer_range
      _lowercase				 =type_vocab_size
      _lowercase				 =layer_norm_eps
      _lowercase				 =use_cache
      _lowercase				 =False







class  lowerCamelCase__  (			lowerCAmelCase):





     @property
     def 	__A   (self			)	->  Mapping[str, Mapping[int, str]]:
      if self.task == "multiple-choice":
       _lowercase				 ={0: '''batch''', 1: '''choice''', 2: '''sequence'''}
      else:
       _lowercase				 ={0: '''batch''', 1: '''sequence'''}
      return OrderedDict(
          [
              ('''input_ids''', dynamic_axis),
              ('''attention_mask''', dynamic_axis),
              ('''token_type_ids''', dynamic_axis),
          ]			)





     @property
     def 	__A   (self			)	->  float:
      return 1e-4

