






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 							a_ :


							'''simple docstring'''






							def __init__(		self			,			lowercase_			,			lowercase_=1_3			,			lowercase_=7			,			lowercase_=False			,			lowercase_=True			,			lowercase_=True			,			lowercase_=True			,			lowercase_=9_9			,			lowercase_=3_2			,			lowercase_=5			,			lowercase_=4			,			lowercase_=3_7			,			lowercase_="gelu"			,			lowercase_=0.1			,			lowercase_=0.1			,			lowercase_=5_1_2			,			lowercase_=1_6			,			lowercase_=2			,			lowercase_=0.02			,			lowercase_=3			,			lowercase_=4			,			lowercase_=None			,			lowercase_=0			,			)					->					str:
														'''simple docstring'''

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

							def      _lowercase       (		self   )					->					Union[str, Any]:
														'''simple docstring'''

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

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

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

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

														lowerCAmelCase_										=     self.get_config()

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

							def      _lowercase       (		self   )					->					int:
														'''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      _lowercase       (		self			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_   )					->					Union[str, Any]:
														'''simple docstring'''

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

							def      _lowercase       (		self			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_   )					->					Optional[int]:
														'''simple docstring'''

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

							def      _lowercase       (		self			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_   )					->					List[Any]:
														'''simple docstring'''

														lowerCAmelCase_										=     DPRReader(config=lowercase_   )
														model.to(lowercase_   )
														model.eval()
														lowerCAmelCase_										=     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      _lowercase       (		self   )					->					Optional[int]:
														'''simple docstring'''

														lowerCAmelCase_										=     self.prepare_config_and_inputs()
														(
														    (
														    lowerCAmelCase_
														)						,   (
														    lowerCAmelCase_
														)						,   (
														    lowerCAmelCase_
														)						,   (
														    lowerCAmelCase_
														)						,   (
														    lowerCAmelCase_
														)						,   (
														    lowerCAmelCase_
														)						,   (
														    lowerCAmelCase_
														)						,   
														)										=     config_and_inputs
														lowerCAmelCase_										=     {'input_ids': input_ids}
														return config, inputs_dict




@require_torch
class 							a_ (						a_	,	a_	,	unittest.TestCase		):


							'''simple docstring'''






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

							__a:	Union[str, Any]     =    False
							__a:	str     =    False  # why?
							__a:	List[Any]     =    False
							__a:	int     =    False

							def      _lowercase       (		self   )					->					Dict:
														'''simple docstring'''

														lowerCAmelCase_										=     DPRModelTester(self   )
														lowerCAmelCase_										=     ConfigTester(self			,			config_class=lowercase_			,			hidden_size=3_7   )

							def      _lowercase       (		self   )					->					Tuple:
														'''simple docstring'''

														self.config_tester.run_common_tests()

							def      _lowercase       (		self   )					->					Dict:
														'''simple docstring'''

														lowerCAmelCase_										=     self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_context_encoder(*lowercase_   )

							def      _lowercase       (		self   )					->					Optional[int]:
														'''simple docstring'''

														lowerCAmelCase_										=     self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_question_encoder(*lowercase_   )

							def      _lowercase       (		self   )					->					Dict:
														'''simple docstring'''

														lowerCAmelCase_										=     self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_reader(*lowercase_   )

							def      _lowercase       (		self   )					->					Union[str, Any]:
														'''simple docstring'''

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

														lowerCAmelCase_										=     DPRQuestionEncoder(config=lowercase_   )
														model.to(lowercase_   )
														model.eval()

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

														self.assertIsNotNone(lowercase_   )






							@slow
							def      _lowercase       (		self   )					->					str:
														'''simple docstring'''

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

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

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

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




@require_torch
class 							a_ (						unittest.TestCase		):


							'''simple docstring'''






							@slow
							def      _lowercase       (		self   )					->					Tuple:
														'''simple docstring'''

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

														lowerCAmelCase_										=     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=lowercase_   )  # [CLS] hello, is my dog cute? [SEP]
														lowerCAmelCase_										=     model(lowercase_   )[0]  # embedding shape = (1, 768)
														# compare the actual values for a slice.
														lowerCAmelCase_										=     torch.tensor(
														    [
														        [
														            0.03_23_62_53,
														            0.12_75_33_35,
														            0.16_81_85_09,
														            0.00_27_97_86,
														            0.3_89_69_33,
														            0.24_26_49_45,
														            0.2_17_89_71,
														            -0.02_33_52_27,
														            -0.08_48_19_59,
														            -0.14_32_41_17,
														        ]
														    ]			,			dtype=torch.float			,			device=lowercase_			,			)
														self.assertTrue(torch.allclose(output[:, :1_0]			,			lowercase_			,			atol=1e-4   )   )






							@slow
							def      _lowercase       (		self   )					->					Optional[int]:
														'''simple docstring'''

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

														lowerCAmelCase_										=     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_   )

														lowerCAmelCase_										=     model(**lowercase_   )

														# compare the actual values for a slice.
														lowerCAmelCase_										=     torch.tensor(
														    [[-10.30_05, -10.77_65, -11.48_72, -11.68_41, -11.93_12, -10.30_02, -9.85_44, -11.73_78, -12.08_21, -10.29_75]]			,			dtype=torch.float			,			device=lowercase_			,			)

														lowerCAmelCase_										=     torch.tensor(
														    [[-11.06_84, -11.70_41, -11.53_97, -10.34_65, -10.87_91, -6.84_43, -11.99_59, -11.03_64, -10.00_96, -6.84_05]]			,			dtype=torch.float			,			device=lowercase_			,			)
														self.assertTrue(torch.allclose(outputs.start_logits[:, :1_0]			,			lowercase_			,			atol=1e-4   )   )
														self.assertTrue(torch.allclose(outputs.end_logits[:, :1_0]			,			lowercase_			,			atol=1e-4   )   )











from dataclasses import dataclass
from typing import Callable, Optional

import datasets




@dataclass
class 							a_ (						datasets.BuilderConfig		):


							'''simple docstring'''






							__a:	Optional[Callable]     =    None
							__a:	Optional[dict]     =    None
							__a:	Optional[datasets.Features]     =    None

							def      _lowercase       (		self   )					->					Optional[int]:
														'''simple docstring'''

														assert self.generator is not None, "generator must be specified"

														if self.gen_kwargs is None:
																					lowerCAmelCase_										=     {}




class 							a_ (						datasets.GeneratorBasedBuilder		):


							'''simple docstring'''






							__a:	Optional[Any]     =    GeneratorConfig

							def      _lowercase       (		self   )					->					Union[str, Any]:
														'''simple docstring'''

														return datasets.DatasetInfo(features=self.config.features   )

							def      _lowercase       (		self			,			lowercase_   )					->					Dict:
														'''simple docstring'''

														return [datasets.SplitGenerator(name=datasets.Split.TRAIN			,			gen_kwargs=self.config.gen_kwargs   )]






							def      _lowercase       (		self			,			**lowercase_   )					->					Union[str, Any]:
														'''simple docstring'''

														for idx, ex in enumerate(self.config.generator(**lowercase_   )   ):
																					yield idx, ex



