


"""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 						snake_case_:






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






     def    lowerCamelCase__   (		self     :    Optional[Any] ):
           lowerCAmelCase    :		int				=      ids_tensor([self.batch_size, self.seq_length]    , self.vocab_size )

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

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

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

           lowerCAmelCase    :		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    , )
           lowerCAmelCase    :		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    lowerCamelCase__   (		self     :    List[str]    , UpperCamelCase_     :    str    , UpperCamelCase_     :    Dict    , UpperCamelCase_     :    str    , UpperCamelCase_     :    Optional[Any]    , UpperCamelCase_     :    str    , UpperCamelCase_     :    Optional[Any]    , UpperCamelCase_     :    List[str] ):
           lowerCAmelCase    :		List[Any]				=      TFDPRContextEncoder(config=_UpperCAmelCase )
           lowerCAmelCase    :		Optional[Any]				=      model(_UpperCAmelCase    , attention_mask=_UpperCAmelCase    , token_type_ids=_UpperCAmelCase )
           lowerCAmelCase    :		Optional[Any]				=      model(_UpperCAmelCase    , token_type_ids=_UpperCAmelCase )
           lowerCAmelCase    :		Optional[Any]				=      model(_UpperCAmelCase )
           self.parent.assertEqual(result.pooler_output.shape    , (self.batch_size, self.projection_dim or self.hidden_size) )






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






     def    lowerCamelCase__   (		self     :    Union[str, Any]    , UpperCamelCase_     :    Optional[Any]    , UpperCamelCase_     :    Optional[Any]    , UpperCamelCase_     :    Dict    , UpperCamelCase_     :    List[str]    , UpperCamelCase_     :    Dict    , UpperCamelCase_     :    List[str]    , UpperCamelCase_     :    List[str] ):
           lowerCAmelCase    :		Union[str, Any]				=      TFDPRReader(config=_UpperCAmelCase )
           lowerCAmelCase    :		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    lowerCamelCase__   (		self     :    Optional[Any] ):
           lowerCAmelCase    :		List[str]				=      self.prepare_config_and_inputs()
           (
               lowerCAmelCase
           )    :		int				=      config_and_inputs
           lowerCAmelCase    :		List[str]				=      {'input_ids': input_ids}
           return config, inputs_dict



@require_tf
class 						snake_case_(				a__							,    a__							,    unittest.TestCase	):
     __UpperCamelCase     =    (
         (
             TFDPRContextEncoder,
             TFDPRQuestionEncoder,
             TFDPRReader,
         )
         if is_tf_available()
         else ()
     )
     __UpperCamelCase     =    {'''feature-extraction''': TFDPRQuestionEncoder} if is_tf_available() else {}

     __UpperCamelCase     =    False
     __UpperCamelCase     =    False
     __UpperCamelCase     =    False
     __UpperCamelCase     =    False
     __UpperCamelCase     =    False






     def    lowerCamelCase__   (		self     :    Tuple ):
           lowerCAmelCase    :		int				=      TFDPRModelTester(self )
           lowerCAmelCase    :		List[str]				=      ConfigTester(self    , config_class=_UpperCAmelCase    , hidden_size=3_7 )






     def    lowerCamelCase__   (		self     :    Union[str, Any] ):
           self.config_tester.run_common_tests()






     def    lowerCamelCase__   (		self     :    Union[str, Any] ):
           lowerCAmelCase    :		Optional[int]				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_context_encoder(*_UpperCAmelCase )






     def    lowerCamelCase__   (		self     :    List[str] ):
           lowerCAmelCase    :		List[str]				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_question_encoder(*_UpperCAmelCase )






     def    lowerCamelCase__   (		self     :    Dict ):
           lowerCAmelCase    :		Tuple				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_dpr_reader(*_UpperCAmelCase )






     @slow
     def    lowerCamelCase__   (		self     :    Dict ):






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

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

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

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



@require_tf
class 						snake_case_(				unittest.TestCase	):






     @slow
     def    lowerCamelCase__   (		self     :    Any ):
           lowerCAmelCase    :		Dict				=      TFDPRQuestionEncoder.from_pretrained('''facebook/dpr-question_encoder-single-nq-base''' )

           lowerCAmelCase    :		List[str]				=      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]
           lowerCAmelCase    :		Optional[int]				=      model(_UpperCAmelCase )[0]  # embedding shape = (1, 768)
           # compare the actual values for a slice.
           lowerCAmelCase    :		Optional[int]				=      tf.constant(
               [
                   [
                       0.03_236_253,
                       0.12_753_335,
                       0.16_818_509,
                       0.00_279_786,
                       0.3_896_933,
                       0.24_264_945,
                       0.2_178_971,
                       -0.02_335_227,
                       -0.08_481_959,
                       -0.14_324_117,
                   ]
               ] )
           self.assertTrue(numpy.allclose(output[:, :1_0].numpy()    , expected_slice.numpy()    , atol=1E-4 ) )









'''simple docstring'''




from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase__    :      List[Any]						=    logging.get_logger(__name__)

lowercase__    :      str						=    {}


class  SCREAMING_SNAKE_CASE    (a__							):
     lowerCAmelCase					=       '''llama'''
     lowerCAmelCase					=       ['''past_key_values''']


     def __init__(						self						,	_UpperCAmelCase=3_2000						,	_UpperCAmelCase=4096						,	_UpperCAmelCase=1_1008						,	_UpperCAmelCase=32						,	_UpperCAmelCase=32						,	_UpperCAmelCase=None						,	_UpperCAmelCase="silu"						,	_UpperCAmelCase=2048						,	_UpperCAmelCase=0.02						,	_UpperCAmelCase=1e-6						,	_UpperCAmelCase=True						,	_UpperCAmelCase=0						,	_UpperCAmelCase=1						,	_UpperCAmelCase=2						,	_UpperCAmelCase=1						,	_UpperCAmelCase=False						,	_UpperCAmelCase=None						,	**_UpperCAmelCase						,	):

          '''simple docstring'''






          __A						: Union[str, Any]      =       vocab_size
          __A						: Union[str, Any]      =       max_position_embeddings
          __A						: Any      =       hidden_size
          __A						: Optional[Any]      =       intermediate_size
          __A						: str      =       num_hidden_layers
          __A						: Optional[Any]      =       num_attention_heads

          # for backward compatibility
          if num_key_value_heads is None:
               __A						: List[Any]      =       num_attention_heads

          __A						: int      =       num_key_value_heads
          __A						: List[Any]      =       hidden_act
          __A						: Union[str, Any]      =       initializer_range
          __A						: List[Any]      =       rms_norm_eps
          __A						: Any      =       pretraining_tp
          __A						: Optional[Any]      =       use_cache
          __A						: Dict      =       rope_scaling
          self._rope_scaling_validation()

          super().__init__(
              pad_token_id=_UpperCAmelCase						,	bos_token_id=_UpperCAmelCase						,	eos_token_id=_UpperCAmelCase						,	tie_word_embeddings=_UpperCAmelCase						,	**_UpperCAmelCase						,	)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          if self.rope_scaling is None:
               return

          if not isinstance(self.rope_scaling						,	_UpperCAmelCase) or len(self.rope_scaling) != 2:
               raise ValueError(
                   '`rope_scaling` must be a dictionary with with two fields, `name` and `factor`, '
                   F'got {self.rope_scaling}')
          __A						: Optional[Any]      =       self.rope_scaling.get('type'						,	_UpperCAmelCase)
          __A						: Tuple      =       self.rope_scaling.get('factor'						,	_UpperCAmelCase)
          if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
               raise ValueError(
                   F'`rope_scaling`\'s name field must be one of [\'linear\', \'dynamic\'], got {rope_scaling_type}')
          if rope_scaling_factor is None or not isinstance(_UpperCAmelCase						,	_UpperCAmelCase) or rope_scaling_factor <= 1.0:
               raise ValueError(F'`rope_scaling`\'s factor field must be an float > 1, got {rope_scaling_factor}')