





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






       '''simple docstring'''


       def __init__(   self     ,    __lowerCamelCase     ,    __lowerCamelCase=1_3     ,    __lowerCamelCase=7     ,    __lowerCamelCase=True     ,    __lowerCamelCase=True     ,    __lowerCamelCase=True     ,    __lowerCamelCase=True     ,    __lowerCamelCase=9_9     ,    __lowerCamelCase=3_2     ,    __lowerCamelCase=2     ,    __lowerCamelCase=4     ,    __lowerCamelCase=3_7     ,    __lowerCamelCase="gelu"     ,    __lowerCamelCase=0.1     ,    __lowerCamelCase=0.1     ,    __lowerCamelCase=5_1_2     ,    __lowerCamelCase=1_6     ,    __lowerCamelCase=2     ,    __lowerCamelCase=0.02     ,    __lowerCamelCase=3     ,    __lowerCamelCase=4     ,    __lowerCamelCase=None     ,    __lowerCamelCase=0     ,    )      ->  str:
            _SCREAMING_SNAKE_CASE							:						List[Any]         =				parent
            _SCREAMING_SNAKE_CASE							:						Optional[int]         =				batch_size
            _SCREAMING_SNAKE_CASE							:						int         =				seq_length
            _SCREAMING_SNAKE_CASE							:						List[Any]         =				is_training
            _SCREAMING_SNAKE_CASE							:						str         =				use_input_mask
            _SCREAMING_SNAKE_CASE							:						str         =				use_token_type_ids
            _SCREAMING_SNAKE_CASE							:						Any         =				use_labels
            _SCREAMING_SNAKE_CASE							:						int         =				vocab_size
            _SCREAMING_SNAKE_CASE							:						int         =				hidden_size
            _SCREAMING_SNAKE_CASE							:						List[str]         =				num_hidden_layers
            _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				num_attention_heads
            _SCREAMING_SNAKE_CASE							:						int         =				intermediate_size
            _SCREAMING_SNAKE_CASE							:						Any         =				hidden_act
            _SCREAMING_SNAKE_CASE							:						Optional[int]         =				hidden_dropout_prob
            _SCREAMING_SNAKE_CASE							:						List[str]         =				attention_probs_dropout_prob
            _SCREAMING_SNAKE_CASE							:						Optional[int]         =				max_position_embeddings
            _SCREAMING_SNAKE_CASE							:						Optional[int]         =				type_vocab_size
            _SCREAMING_SNAKE_CASE							:						List[str]         =				type_sequence_label_size
            _SCREAMING_SNAKE_CASE							:						int         =				initializer_range
            _SCREAMING_SNAKE_CASE							:						Dict         =				num_labels
            _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				num_choices
            _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				scope
            _SCREAMING_SNAKE_CASE							:						List[Any]         =				projection_dim




       def      UpperCamelCase_      (   self				)      ->  Tuple:
            _SCREAMING_SNAKE_CASE							:						List[Any]         =				ids_tensor([self.batch_size, self.seq_length]     ,    self.vocab_size				)

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

            _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				None
            if self.use_token_type_ids:
                 _SCREAMING_SNAKE_CASE							:						Optional[int]         =				ids_tensor([self.batch_size, self.seq_length]     ,    self.type_vocab_size				)

            _SCREAMING_SNAKE_CASE							:						int         =				None
            _SCREAMING_SNAKE_CASE							:						List[Any]         =				None
            _SCREAMING_SNAKE_CASE							:						int         =				None
            if self.use_labels:
                 _SCREAMING_SNAKE_CASE							:						Dict         =				ids_tensor([self.batch_size]     ,    self.type_sequence_label_size				)
                 _SCREAMING_SNAKE_CASE							:						Union[str, Any]         =				ids_tensor([self.batch_size, self.seq_length]     ,    self.num_labels				)
                 _SCREAMING_SNAKE_CASE							:						Tuple         =				ids_tensor([self.batch_size]     ,    self.num_choices				)

            _SCREAMING_SNAKE_CASE							:						Optional[int]         =				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=_lowerCamelCase     ,    initializer_range=self.initializer_range     ,    )
            _SCREAMING_SNAKE_CASE							:						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      UpperCamelCase_      (   self     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase				)      ->  Union[str, Any]:
            _SCREAMING_SNAKE_CASE							:						Tuple         =				TFDPRContextEncoder(config=_lowerCamelCase				)
            _SCREAMING_SNAKE_CASE							:						Dict         =				model(_lowerCamelCase     ,    attention_mask=_lowerCamelCase     ,    token_type_ids=_lowerCamelCase				)
            _SCREAMING_SNAKE_CASE							:						Dict         =				model(_lowerCamelCase     ,    token_type_ids=_lowerCamelCase				)
            _SCREAMING_SNAKE_CASE							:						int         =				model(_lowerCamelCase				)
            self.parent.assertEqual(result.pooler_output.shape     ,    (self.batch_size, self.projection_dim or self.hidden_size)				)




       def      UpperCamelCase_      (   self     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase				)      ->  str:
            _SCREAMING_SNAKE_CASE							:						List[str]         =				TFDPRQuestionEncoder(config=_lowerCamelCase				)
            _SCREAMING_SNAKE_CASE							:						List[Any]         =				model(_lowerCamelCase     ,    attention_mask=_lowerCamelCase     ,    token_type_ids=_lowerCamelCase				)
            _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				model(_lowerCamelCase     ,    token_type_ids=_lowerCamelCase				)
            _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				model(_lowerCamelCase				)
            self.parent.assertEqual(result.pooler_output.shape     ,    (self.batch_size, self.projection_dim or self.hidden_size)				)




       def      UpperCamelCase_      (   self     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase				)      ->  List[Any]:
            _SCREAMING_SNAKE_CASE							:						List[Any]         =				TFDPRReader(config=_lowerCamelCase				)
            _SCREAMING_SNAKE_CASE							:						Tuple         =				model(_lowerCamelCase     ,    attention_mask=_lowerCamelCase				)

            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      UpperCamelCase_      (   self				)      ->  Any:
            _SCREAMING_SNAKE_CASE							:						Any         =				self.prepare_config_and_inputs()
            (
                _SCREAMING_SNAKE_CASE							
            )							:						Dict         =				config_and_inputs
            _SCREAMING_SNAKE_CASE							:						Optional[int]         =				{'''input_ids''': input_ids}
            return config, inputs_dict





@require_tf
class       lowerCAmelCase__(			_snake_case	,  _snake_case	,  unittest.TestCase     ):






       '''simple docstring'''


       __snake_case			      =     (
           (
               TFDPRContextEncoder,
               TFDPRQuestionEncoder,
               TFDPRReader,
           )
           if is_tf_available()
           else ()
       )
       __snake_case			      =     {'feature-extraction': TFDPRQuestionEncoder} if is_tf_available() else {}

       __snake_case			      =     False
       __snake_case			      =     False
       __snake_case			      =     False
       __snake_case			      =     False
       __snake_case			      =     False




       def      UpperCamelCase_      (   self				)      ->  str:
            _SCREAMING_SNAKE_CASE							:						int         =				TFDPRModelTester(self				)
            _SCREAMING_SNAKE_CASE							:						Dict         =				ConfigTester(self     ,    config_class=_lowerCamelCase     ,    hidden_size=3_7				)




       def      UpperCamelCase_      (   self				)      ->  Any:
            self.config_tester.run_common_tests()




       def      UpperCamelCase_      (   self				)      ->  List[Any]:
            _SCREAMING_SNAKE_CASE							:						Any         =				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_dpr_context_encoder(*_lowerCamelCase				)




       def      UpperCamelCase_      (   self				)      ->  List[Any]:
            _SCREAMING_SNAKE_CASE							:						str         =				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_dpr_question_encoder(*_lowerCamelCase				)




       def      UpperCamelCase_      (   self				)      ->  Optional[Any]:
            _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_dpr_reader(*_lowerCamelCase				)







       @slow
       def      UpperCamelCase_      (   self				)      ->  Union[str, Any]:
            for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 _SCREAMING_SNAKE_CASE							:						Any         =				TFDPRContextEncoder.from_pretrained(_lowerCamelCase				)
                 self.assertIsNotNone(_lowerCamelCase				)

            for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 _SCREAMING_SNAKE_CASE							:						str         =				TFDPRContextEncoder.from_pretrained(_lowerCamelCase				)
                 self.assertIsNotNone(_lowerCamelCase				)

            for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				TFDPRQuestionEncoder.from_pretrained(_lowerCamelCase				)
                 self.assertIsNotNone(_lowerCamelCase				)

            for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 _SCREAMING_SNAKE_CASE							:						Tuple         =				TFDPRReader.from_pretrained(_lowerCamelCase				)
                 self.assertIsNotNone(_lowerCamelCase				)





@require_tf
class       lowerCAmelCase__(			unittest.TestCase     ):






       '''simple docstring'''


       @slow
       def      UpperCamelCase_      (   self				)      ->  str:
            _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				TFDPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base"				)

            _SCREAMING_SNAKE_CASE							:						Optional[Any]         =				tf.constant(
                [[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]]				)  # [CLS] hello, is my dog cute? [SEP]
            _SCREAMING_SNAKE_CASE							:						List[str]         =				model(_lowerCamelCase				)[0]  # embedding shape = (1, 768)
            # compare the actual values for a slice.
            _SCREAMING_SNAKE_CASE							:						List[str]         =				tf.constant(
                [
                    [
                        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,
                    ]
                ]				)
            self.assertTrue(numpy.allclose(output[:, :1_0].numpy()     ,    expected_slice.numpy()     ,    atol=1E-4				)				)



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available


snake_case							:       Any										=	{
    '''configuration_gpt_neo''': ['''GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''GPTNeoConfig''', '''GPTNeoOnnxConfig'''],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case							:       List[Any]										=	[
         '''GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''GPTNeoForCausalLM''',
         '''GPTNeoForQuestionAnswering''',
         '''GPTNeoForSequenceClassification''',
         '''GPTNeoForTokenClassification''',
         '''GPTNeoModel''',
         '''GPTNeoPreTrainedModel''',
         '''load_tf_weights_in_gpt_neo''',
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case							:       Any										=	[
         '''FlaxGPTNeoForCausalLM''',
         '''FlaxGPTNeoModel''',
         '''FlaxGPTNeoPreTrainedModel''',
     ]


if TYPE_CHECKING:
     from .configuration_gpt_neo import GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoConfig, GPTNeoOnnxConfig

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_gpt_neo import (
              GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST,
              GPTNeoForCausalLM,
              GPTNeoForQuestionAnswering,
              GPTNeoForSequenceClassification,
              GPTNeoForTokenClassification,
              GPTNeoModel,
              GPTNeoPreTrainedModel,
              load_tf_weights_in_gpt_neo,
          )

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel, FlaxGPTNeoPreTrainedModel


else:
     import sys

     snake_case							:       Optional[int]										=	_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





