from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import numpy
   import tensorflow as tf

   from transformers import (
       TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
       TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
       TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
       BertConfig,
       DPRConfig,
       TFDPRContextEncoder,
       TFDPRQuestionEncoder,
       TFDPRReader,
   )







class __snake_case				:




   def __init__(    self			:				Tuple       ,	_lowercase			:				Optional[Any]       ,	_lowercase			:				Dict=13       ,	_lowercase			:				Optional[int]=7       ,	_lowercase			:				int=True       ,	_lowercase			:				List[str]=True       ,	_lowercase			:				str=True       ,	_lowercase			:				Optional[Any]=True       ,	_lowercase			:				int=99       ,	_lowercase			:				Tuple=32       ,	_lowercase			:				Dict=2       ,	_lowercase			:				Any=4       ,	_lowercase			:				List[Any]=37       ,	_lowercase			:				Any="gelu"       ,	_lowercase			:				Any=0.1       ,	_lowercase			:				Tuple=0.1       ,	_lowercase			:				Optional[int]=5_12       ,	_lowercase			:				int=16       ,	_lowercase			:				str=2       ,	_lowercase			:				int=0.02       ,	_lowercase			:				Dict=3       ,	_lowercase			:				Any=4       ,	_lowercase			:				List[Any]=None       ,	_lowercase			:				Optional[Any]=0       ,	):



     """simple docstring"""






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




   def 	__a						(    self			:				Tuple      ):



     """simple docstring"""






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

     SCREAMING_SNAKE_CASE__										=      None
     if self.use_input_mask:
       # follow test_modeling_tf_ctrl.py
       SCREAMING_SNAKE_CASE__										=      random_attention_mask([self.batch_size, self.seq_length]      )

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

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

     SCREAMING_SNAKE_CASE__										=      BertConfig(
         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       ,	is_decoder=_lowercase       ,	initializer_range=self.initializer_range       ,	)
     SCREAMING_SNAKE_CASE__										=      DPRConfig(projection_dim=self.projection_dim       ,	**config.to_dict()      )

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




   def 	__a						(    self			:				Any       ,	_lowercase			:				str       ,	_lowercase			:				Optional[int]       ,	_lowercase			:				List[str]       ,	_lowercase			:				Optional[int]       ,	_lowercase			:				Union[str, Any]       ,	_lowercase			:				List[str]       ,	_lowercase			:				List[Any]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      TFDPRContextEncoder(config=_lowercase      )
     SCREAMING_SNAKE_CASE__										=      model(_lowercase       ,	attention_mask=_lowercase       ,	token_type_ids=_lowercase      )
     SCREAMING_SNAKE_CASE__										=      model(_lowercase       ,	token_type_ids=_lowercase      )
     SCREAMING_SNAKE_CASE__										=      model(_lowercase      )
     self.parent.assertEqual(result.pooler_output.shape       ,	(self.batch_size, self.projection_dim or self.hidden_size)      )




   def 	__a						(    self			:				int       ,	_lowercase			:				Dict       ,	_lowercase			:				Any       ,	_lowercase			:				List[Any]       ,	_lowercase			:				List[str]       ,	_lowercase			:				List[str]       ,	_lowercase			:				Optional[int]       ,	_lowercase			:				int      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      TFDPRQuestionEncoder(config=_lowercase      )
     SCREAMING_SNAKE_CASE__										=      model(_lowercase       ,	attention_mask=_lowercase       ,	token_type_ids=_lowercase      )
     SCREAMING_SNAKE_CASE__										=      model(_lowercase       ,	token_type_ids=_lowercase      )
     SCREAMING_SNAKE_CASE__										=      model(_lowercase      )
     self.parent.assertEqual(result.pooler_output.shape       ,	(self.batch_size, self.projection_dim or self.hidden_size)      )




   def 	__a						(    self			:				Any       ,	_lowercase			:				Any       ,	_lowercase			:				Union[str, Any]       ,	_lowercase			:				str       ,	_lowercase			:				int       ,	_lowercase			:				Optional[Any]       ,	_lowercase			:				str       ,	_lowercase			:				Optional[Any]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      TFDPRReader(config=_lowercase      )
     SCREAMING_SNAKE_CASE__										=      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 	__a						(    self			:				Optional[int]      ):



     """simple docstring"""






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







@require_tf
class __snake_case				(						snake_case__   ,		snake_case__   ,		unittest.TestCase		):
   lowerCAmelCase_														=							(
       (
           TFDPRContextEncoder,
           TFDPRQuestionEncoder,
           TFDPRReader,
       )
       if is_tf_available()
       else ()
   )
   lowerCAmelCase_														=							{"feature-extraction": TFDPRQuestionEncoder} if is_tf_available() else {}

   lowerCAmelCase_														=							False
   lowerCAmelCase_														=							False
   lowerCAmelCase_														=							False
   lowerCAmelCase_														=							False
   lowerCAmelCase_														=							False




   def 	__a						(    self			:				str      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      TFDPRModelTester(self      )
     SCREAMING_SNAKE_CASE__										=      ConfigTester(self       ,	config_class=_lowercase       ,	hidden_size=37      )




   def 	__a						(    self			:				Any      ):



     """simple docstring"""






     self.config_tester.run_common_tests()




   def 	__a						(    self			:				Any      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_dpr_context_encoder(*_lowercase      )




   def 	__a						(    self			:				Dict      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_dpr_question_encoder(*_lowercase      )




   def 	__a						(    self			:				int      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_dpr_reader(*_lowercase      )




   @slow
   def 	__a						(    self			:				Any      ):



     """simple docstring"""






     for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       SCREAMING_SNAKE_CASE__										=      TFDPRContextEncoder.from_pretrained(_lowercase      )
       self.assertIsNotNone(_lowercase      )

     for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       SCREAMING_SNAKE_CASE__										=      TFDPRContextEncoder.from_pretrained(_lowercase      )
       self.assertIsNotNone(_lowercase      )

     for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       SCREAMING_SNAKE_CASE__										=      TFDPRQuestionEncoder.from_pretrained(_lowercase      )
       self.assertIsNotNone(_lowercase      )

     for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       SCREAMING_SNAKE_CASE__										=      TFDPRReader.from_pretrained(_lowercase      )
       self.assertIsNotNone(_lowercase      )







@require_tf
class __snake_case				(						unittest.TestCase		):




   @slow
   def 	__a						(    self			:				Dict      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      TFDPRQuestionEncoder.from_pretrained("""facebook/dpr-question_encoder-single-nq-base"""      )

     SCREAMING_SNAKE_CASE__										=      tf.constant(
         [[1_01, 75_92, 10_10, 20_03, 20_26, 38_99, 1_01_40, 10_29, 1_02]]      )  # [CLS] hello, is my dog cute? [SEP]
     SCREAMING_SNAKE_CASE__										=      model(_lowercase      )[0]  # embedding shape = (1, 768)
     # compare the actual values for a slice.
     SCREAMING_SNAKE_CASE__										=      tf.constant(
         [
             [
                 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,
             ]
         ]      )
     self.assertTrue(numpy.allclose(output[:, :10].numpy()       ,	expected_slice.numpy()       ,	atol=1E-4      )      )








import enum
import os
from hashlib import shaaaa
from typing import Optional

from .. import config
from .logging import get_logger


_a      	=  get_logger(__name__)





class __lowerCamelCase (      enum.Enum):





	"""simple docstring"""






	UpperCamelCase__								=      "all_checks"
	UpperCamelCase__								=      "basic_checks"
	UpperCamelCase__								=      "no_checks"





class __lowerCamelCase (      snake_case__):





	"""simple docstring"""











class __lowerCamelCase (      snake_case__):





	"""simple docstring"""











class __lowerCamelCase (      snake_case__):





	"""simple docstring"""











class __lowerCamelCase (      snake_case__):





	"""simple docstring"""









def __A    (      __lowerCAmelCase ,		__lowerCAmelCase ,		__lowerCAmelCase=None	)->  str:





						"""simple docstring"""



						if expected_checksums is None:
												logger.info('Unable to verify checksums.'	)
												return
						if len(set(__lowerCAmelCase	) - set(__lowerCAmelCase	)	) > 0:
												raise ExpectedMoreDownloadedFiles(str(set(__lowerCAmelCase	) - set(__lowerCAmelCase	)	)	)
						if len(set(__lowerCAmelCase	) - set(__lowerCAmelCase	)	) > 0:
												raise UnexpectedDownloadedFile(str(set(__lowerCAmelCase	) - set(__lowerCAmelCase	)	)	)
						_UpperCAmelCase									=			[url for url in expected_checksums if expected_checksums[url] != recorded_checksums[url]]
						_UpperCAmelCase									=			' for ' + verification_name if verification_name is not None else ''
						if len(__lowerCAmelCase	) > 0:
												raise NonMatchingChecksumError(
												    F"""Checksums didn't match{for_verification_name}:\n"""
												    F"""{bad_urls}\n"""
												    'Set `verification_mode=\'no_checks\'` to skip checksums verification and ignore this error'	)
						logger.info('All the checksums matched successfully' + for_verification_name	)





class __lowerCamelCase (      snake_case__):





	"""simple docstring"""











class __lowerCamelCase (      snake_case__):





	"""simple docstring"""











class __lowerCamelCase (      snake_case__):





	"""simple docstring"""











class __lowerCamelCase (      snake_case__):





	"""simple docstring"""









def __A    (      __lowerCAmelCase ,		__lowerCAmelCase	)->  int:





						"""simple docstring"""



						if expected_splits is None:
												logger.info('Unable to verify splits sizes.'	)
												return
						if len(set(__lowerCAmelCase	) - set(__lowerCAmelCase	)	) > 0:
												raise ExpectedMoreSplits(str(set(__lowerCAmelCase	) - set(__lowerCAmelCase	)	)	)
						if len(set(__lowerCAmelCase	) - set(__lowerCAmelCase	)	) > 0:
												raise UnexpectedSplits(str(set(__lowerCAmelCase	) - set(__lowerCAmelCase	)	)	)
						_UpperCAmelCase									=			[
						    {'expected': expected_splits[name], 'recorded': recorded_splits[name]}
						    for name in expected_splits
						    if expected_splits[name].num_examples != recorded_splits[name].num_examples
						]
						if len(__lowerCAmelCase	) > 0:
												raise NonMatchingSplitsSizesError(str(__lowerCAmelCase	)	)
						logger.info('All the splits matched successfully.'	)


def __A    (      __lowerCAmelCase ,		__lowerCAmelCase = True	)->  dict:





						"""simple docstring"""



						if record_checksum:
												_UpperCAmelCase									=			shaaaa()
												with open(__lowerCAmelCase ,		'rb'	) as f:
																		for chunk in iter(lambda: f.read(1 << 20	) ,		b''	):
																								m.update(__lowerCAmelCase	)
																		_UpperCAmelCase									=			m.hexdigest()
						else:
												_UpperCAmelCase									=			None
						return {"num_bytes": os.path.getsize(__lowerCAmelCase	), "checksum": checksum}


def __A    (      __lowerCAmelCase	)->  List[str]:





						"""simple docstring"""



						if dataset_size and config.IN_MEMORY_MAX_SIZE:
												return dataset_size < config.IN_MEMORY_MAX_SIZE
						else:
												return False
