




'''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 __magic_name__:






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



									'''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						:						List[Any]       ):



									'''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						:						List[Any]       ):



									'''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						:						Optional[int]					,   __UpperCamelCase						:						Any					,   __UpperCamelCase						:						Optional[int]					,   __UpperCamelCase						:						Union[str, Any]					,   __UpperCamelCase						:						List[str]					,   __UpperCamelCase						:						int					,   __UpperCamelCase						:						Tuple					,   __UpperCamelCase						:						List[Any]       ):



									'''simple docstring'''






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






						def       __lowerCAmelCase(   self						:						Optional[int]					,   __UpperCamelCase						:						Any					,   __UpperCamelCase						:						str					,   __UpperCamelCase						:						List[str]					,   __UpperCamelCase						:						Dict					,   __UpperCamelCase						:						int					,   __UpperCamelCase						:						str					,   __UpperCamelCase						:						List[Any]       ):



									'''simple docstring'''






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






						def       __lowerCAmelCase(   self						:						Optional[Any]					,   __UpperCamelCase						:						Optional[int]					,   __UpperCamelCase						:						Optional[int]					,   __UpperCamelCase						:						Union[str, Any]					,   __UpperCamelCase						:						Dict					,   __UpperCamelCase						:						Any					,   __UpperCamelCase						:						Tuple					,   __UpperCamelCase						:						Any       ):



									'''simple docstring'''






									snake_case__          =						DPRReader(config=a__       )
									model.to(a__       )
									model.eval()
									snake_case__          =						model(
									    a__					,   attention_mask=a__					,   )

									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						:						List[Any]       ):



									'''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 __magic_name__(						_UpperCamelCase    ,		_UpperCamelCase    ,		unittest.TestCase      ):
						UpperCAmelCase_		:					Optional[int]    =  (
						    (
						        DPRContextEncoder,
						        DPRQuestionEncoder,
						        DPRReader,
						    )
						    if is_torch_available()
						    else ()
						)
						UpperCAmelCase_		:					Dict    =  {"""feature-extraction""": DPRQuestionEncoder} if is_torch_available() else {}

						UpperCAmelCase_		:					Optional[Any]    =  False
						UpperCAmelCase_		:					Dict    =  False  # why?
						UpperCAmelCase_		:					Dict    =  False
						UpperCAmelCase_		:					Dict    =  False






						def       __lowerCAmelCase(   self						:						Optional[int]       ):



									'''simple docstring'''






									snake_case__          =						DPRModelTester(self       )
									snake_case__          =						ConfigTester(self					,   config_class=a__					,   hidden_size=3_7       )






						def       __lowerCAmelCase(   self						:						Tuple       ):



									'''simple docstring'''






									self.config_tester.run_common_tests()






						def       __lowerCAmelCase(   self						:						Dict       ):



									'''simple docstring'''






									snake_case__          =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_context_encoder(*a__       )






						def       __lowerCAmelCase(   self						:						str       ):



									'''simple docstring'''






									snake_case__          =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_question_encoder(*a__       )






						def       __lowerCAmelCase(   self						:						Dict       ):



									'''simple docstring'''






									snake_case__          =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_reader(*a__       )






						def       __lowerCAmelCase(   self						:						List[str]       ):



									'''simple docstring'''






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

									snake_case__          =						DPRQuestionEncoder(config=a__       )
									model.to(a__       )
									model.eval()

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

									self.assertIsNotNone(a__       )







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



									'''simple docstring'''






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

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

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

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







@require_torch
class __magic_name__(						unittest.TestCase      ):






						@slow
						def       __lowerCAmelCase(   self						:						List[str]       ):



									'''simple docstring'''






									snake_case__          =						DPRQuestionEncoder.from_pretrained("""facebook/dpr-question_encoder-single-nq-base"""					,   return_dict=a__       )
									model.to(a__       )

									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=a__       )  # [CLS] hello, is my dog cute? [SEP]
									snake_case__          =						model(a__       )[0]  # embedding shape = (1, 768)
									# compare the actual values for a slice.
									snake_case__          =						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=a__					,   )
									self.assertTrue(torch.allclose(output[:, :1_0]					,   a__					,   atol=1E-4       )       )







						@slow
						def       __lowerCAmelCase(   self						:						Tuple       ):



									'''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(a__       )

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

									snake_case__          =						model(**a__       )

									# compare the actual values for a slice.
									snake_case__          =						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.85_44, -1_1.7_3_7_8, -1_2.0_8_2_1, -1_0.2_9_7_5]]					,   dtype=torch.float					,   device=a__					,   )

									snake_case__          =						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.84_43, -1_1.9_9_5_9, -1_1.0_3_6_4, -1_0.0_0_9_6, -6.84_05]]					,   dtype=torch.float					,   device=a__					,   )
									self.assertTrue(torch.allclose(outputs.start_logits[:, :1_0]					,   a__					,   atol=1E-4       )       )
									self.assertTrue(torch.allclose(outputs.end_logits[:, :1_0]					,   a__					,   atol=1E-4       )       )



"""simple docstring"""


import os

from apache_beam.io.filesystems import FileSystems
from apache_beam.pipeline import Pipeline

from .logging import get_logger


a_											=     2 << 20  # 2mb
a_											=     get_logger(__name__)
class snake_case   (	_UpperCamelCase):







		def 			a_			(       self   :							Tuple		)							->    Optional[int]:





						'''simple docstring'''







						_A       	=      self._options.get_all_options().get("runner"		)
						return runner in [None, "DirectRunner", "PortableRunner"]






def   a__			(					__lowercase  ,      __lowercase  ,      __lowercase=False )						->						Any:
				_A       	=      FileSystems
				if fs.exists(__lowercase ):
								if force_upload:
												logger.info(f"""Remote path already exist: {remote_file_path}. Overwriting it as force_upload=True.""" )
								else:
												logger.info(f"""Remote path already exist: {remote_file_path}. Skipping it as force_upload=False.""" )
												return
				with fs.create(__lowercase ) as remote_file:
								with open(__lowercase  ,      "rb" ) as local_file:
												_A       	=      local_file.read(__lowercase )
												while chunk:
																remote_file.write(__lowercase )
																_A       	=      local_file.read(__lowercase )






def   a__			(					__lowercase  ,      __lowercase  ,      __lowercase=False )						->						Union[str, Any]:
				_A       	=      FileSystems
				if os.path.exists(__lowercase ):
								if force_download:
												logger.info(f"""Local path already exist: {remote_file_path}. Overwriting it as force_upload=True.""" )
								else:
												logger.info(f"""Local path already exist: {remote_file_path}. Skipping it as force_upload=False.""" )
												return
				with fs.open(__lowercase ) as remote_file:
								with open(__lowercase  ,      "wb" ) as local_file:
												_A       	=      remote_file.read(__lowercase )
												while chunk:
																local_file.write(__lowercase )
																_A       	=      remote_file.read(__lowercase )