





'''simple docstring'''




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


     def __init__(						self						,	_UpperCAmelCase						,	_UpperCAmelCase=13						,	_UpperCAmelCase=7						,	_UpperCAmelCase=True						,	_UpperCAmelCase=True						,	_UpperCAmelCase=True						,	_UpperCAmelCase=True						,	_UpperCAmelCase=99						,	_UpperCAmelCase=32						,	_UpperCAmelCase=2						,	_UpperCAmelCase=4						,	_UpperCAmelCase=37						,	_UpperCAmelCase="gelu"						,	_UpperCAmelCase=0.1						,	_UpperCAmelCase=0.1						,	_UpperCAmelCase=512						,	_UpperCAmelCase=16						,	_UpperCAmelCase=2						,	_UpperCAmelCase=0.02						,	_UpperCAmelCase=3						,	_UpperCAmelCase=4						,	_UpperCAmelCase=None						,	_UpperCAmelCase=0						,	):

          '''simple docstring'''






          __A						: List[Any]      =       parent
          __A						: str      =       batch_size
          __A						: List[Any]      =       seq_length
          __A						: Any      =       is_training
          __A						: List[str]      =       use_input_mask
          __A						: Dict      =       use_token_type_ids
          __A						: Optional[Any]      =       use_labels
          __A						: Tuple      =       vocab_size
          __A						: int      =       hidden_size
          __A						: int      =       num_hidden_layers
          __A						: Any      =       num_attention_heads
          __A						: Optional[int]      =       intermediate_size
          __A						: List[Any]      =       hidden_act
          __A						: str      =       hidden_dropout_prob
          __A						: List[Any]      =       attention_probs_dropout_prob
          __A						: List[Any]      =       max_position_embeddings
          __A						: Any      =       type_vocab_size
          __A						: Any      =       type_sequence_label_size
          __A						: Tuple      =       initializer_range
          __A						: Any      =       num_labels
          __A						: Optional[Any]      =       num_choices
          __A						: int      =       scope
          __A						: Optional[int]      =       projection_dim


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: int      =       ids_tensor([self.batch_size, self.seq_length]						,	self.vocab_size)

          __A						: List[str]      =       None
          if self.use_input_mask:
               # follow test_modeling_tf_ctrl.py
               __A						: Optional[Any]      =       random_attention_mask([self.batch_size, self.seq_length])

          __A						: List[Any]      =       None
          if self.use_token_type_ids:
               __A						: str      =       ids_tensor([self.batch_size, self.seq_length]						,	self.type_vocab_size)

          __A						: Any      =       None
          __A						: List[Any]      =       None
          __A						: Dict      =       None
          if self.use_labels:
               __A						: int      =       ids_tensor([self.batch_size]						,	self.type_sequence_label_size)
               __A						: Optional[int]      =       ids_tensor([self.batch_size, self.seq_length]						,	self.num_labels)
               __A						: Optional[int]      =       ids_tensor([self.batch_size]						,	self.num_choices)

          __A						: Dict      =       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=_UpperCAmelCase						,	initializer_range=self.initializer_range						,	)
          __A						: List[str]      =       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 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase):

          '''simple docstring'''






          __A						: List[Any]      =       TFDPRContextEncoder(config=_UpperCAmelCase)
          __A						: Optional[Any]      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase						,	token_type_ids=_UpperCAmelCase)
          __A						: Optional[Any]      =       model(_UpperCAmelCase						,	token_type_ids=_UpperCAmelCase)
          __A						: Optional[Any]      =       model(_UpperCAmelCase)
          self.parent.assertEqual(result.pooler_output.shape						,	(self.batch_size, self.projection_dim or self.hidden_size))


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase):

          '''simple docstring'''






          __A						: Union[str, Any]      =       TFDPRQuestionEncoder(config=_UpperCAmelCase)
          __A						: List[str]      =       model(_UpperCAmelCase						,	attention_mask=_UpperCAmelCase						,	token_type_ids=_UpperCAmelCase)
          __A						: Optional[int]      =       model(_UpperCAmelCase						,	token_type_ids=_UpperCAmelCase)
          __A						: Optional[int]      =       model(_UpperCAmelCase)
          self.parent.assertEqual(result.pooler_output.shape						,	(self.batch_size, self.projection_dim or self.hidden_size))


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase						,	_UpperCAmelCase):

          '''simple docstring'''






          __A						: Union[str, Any]      =       TFDPRReader(config=_UpperCAmelCase)
          __A						: Dict      =       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 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[str]      =       self.prepare_config_and_inputs()
          (
              (
              __A
          ) ,(
              __A
          ) ,(
              __A
          ) ,(
              __A
          ) ,(
              __A
          ) ,(
              __A
          ) ,(
              __A
          ) ,
          )						: int      =       config_and_inputs
          __A						: List[str]      =       {'input_ids': input_ids}
          return config, inputs_dict




@require_tf
class  SCREAMING_SNAKE_CASE    (a__					,		a__					,		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 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: int      =       TFDPRModelTester(self)
          __A						: List[str]      =       ConfigTester(self						,	config_class=_UpperCAmelCase						,	hidden_size=37)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          self.config_tester.run_common_tests()


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Optional[int]      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_dpr_context_encoder(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[str]      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_dpr_question_encoder(*_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Tuple      =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_dpr_reader(*_UpperCAmelCase)


     @slow
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __A						: Dict      =       TFDPRContextEncoder.from_pretrained(_UpperCAmelCase)
               self.assertIsNotNone(_UpperCAmelCase)

          for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __A						: Dict      =       TFDPRContextEncoder.from_pretrained(_UpperCAmelCase)
               self.assertIsNotNone(_UpperCAmelCase)

          for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __A						: str      =       TFDPRQuestionEncoder.from_pretrained(_UpperCAmelCase)
               self.assertIsNotNone(_UpperCAmelCase)

          for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __A						: str      =       TFDPRReader.from_pretrained(_UpperCAmelCase)
               self.assertIsNotNone(_UpperCAmelCase)




@require_tf
class  SCREAMING_SNAKE_CASE    (unittest.TestCase							):


     @slow
     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Dict      =       TFDPRQuestionEncoder.from_pretrained('facebook/dpr-question_encoder-single-nq-base')

          __A						: List[str]      =       tf.constant(
              [[101, 7592, 1010, 2003, 2026, 3899, 1_0140, 1029, 102]])  # [CLS] hello, is my dog cute? [SEP]
          __A						: Optional[int]      =       model(_UpperCAmelCase)[0]  # embedding shape = (1, 768)
          # compare the actual values for a slice.
          __A						: Optional[int]      =       tf.constant(
              [
                  [
                      0.03236253,
                      0.12753335,
                      0.16818509,
                      0.00279786,
                      0.3896933,
                      0.24264945,
                      0.2178971,
                      -0.02335227,
                      -0.08481959,
                      -0.14324117,
                  ]
              ])
          self.assertTrue(numpy.allclose(output[:, :10].numpy()						,	expected_slice.numpy()						,	atol=1e-4))






'''simple docstring'''




from argparse import ArgumentParser

from datasets.commands.convert import ConvertCommand
from datasets.commands.dummy_data import DummyDataCommand
from datasets.commands.env import EnvironmentCommand
from datasets.commands.run_beam import RunBeamCommand
from datasets.commands.test import TestCommand
from datasets.utils.logging import set_verbosity_info






def        _lowerCAmelCase			(  __snake_case  : Tuple			)     ->     Dict:
     return {key.lstrip('-'			): value for key, value in zip(unknown_args[::2] ,							unknown_args[1::2]			)}






def        _lowerCAmelCase			(  )     ->     Tuple:
     __A						: int      =       ArgumentParser(
         'HuggingFace Datasets CLI tool' ,							usage='datasets-cli <command> [<args>]' ,							allow_abbrev=__snake_case			)
     __A						: Optional[Any]      =       parser.add_subparsers(help='datasets-cli command helpers'			)
     set_verbosity_info()

     # Register commands
     ConvertCommand.register_subcommand(__snake_case			)
     EnvironmentCommand.register_subcommand(__snake_case			)
     TestCommand.register_subcommand(__snake_case			)
     RunBeamCommand.register_subcommand(__snake_case			)
     DummyDataCommand.register_subcommand(__snake_case			)

     # Parse args
     __A ,__A						: Optional[Any]      =       parser.parse_known_args()
     if not hasattr(__snake_case ,							'func'			):
          parser.print_help()
          exit(1			)
     __A						: Any      =       parse_unknown_args(__snake_case			)

     # Run
     __A						: List[Any]      =       args.func(__snake_case ,							**__snake_case			)
     service.run()


if __name__ == "__main__":
  main()