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





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




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : List[str]	     =       ids_tensor([self.batch_size, self.seq_length]	,							self.vocab_size       )

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

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

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

       UpperCamelCase     : Optional[Any]	     =       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=A_	,							initializer_range=self.initializer_range	,							)
       UpperCamelCase     : 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 	__UpperCamelCase(			self	,							A_	,							A_	,							A_	,							A_	,							A_	,							A_	,							A_       ):
       '''simple docstring'''




       UpperCamelCase     : int	     =       TFDPRContextEncoder(config=A_       )
       UpperCamelCase     : Union[str, Any]	     =       model(A_	,							attention_mask=A_	,							token_type_ids=A_       )
       UpperCamelCase     : Union[str, Any]	     =       model(A_	,							token_type_ids=A_       )
       UpperCamelCase     : List[Any]	     =       model(A_       )
       self.parent.assertEqual(result.pooler_output.shape	,							(self.batch_size, self.projection_dim or self.hidden_size)       )





      def 	__UpperCamelCase(			self	,							A_	,							A_	,							A_	,							A_	,							A_	,							A_	,							A_       ):
       '''simple docstring'''




       UpperCamelCase     : List[str]	     =       TFDPRQuestionEncoder(config=A_       )
       UpperCamelCase     : Dict	     =       model(A_	,							attention_mask=A_	,							token_type_ids=A_       )
       UpperCamelCase     : Optional[Any]	     =       model(A_	,							token_type_ids=A_       )
       UpperCamelCase     : List[str]	     =       model(A_       )
       self.parent.assertEqual(result.pooler_output.shape	,							(self.batch_size, self.projection_dim or self.hidden_size)       )





      def 	__UpperCamelCase(			self	,							A_	,							A_	,							A_	,							A_	,							A_	,							A_	,							A_       ):
       '''simple docstring'''




       UpperCamelCase     : List[Any]	     =       TFDPRReader(config=A_       )
       UpperCamelCase     : int	     =       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 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : List[str]	     =       self.prepare_config_and_inputs()
       (
           UpperCamelCase  
       )     : Optional[int]	     =       config_and_inputs
       UpperCamelCase     : Dict	     =       {"input_ids": input_ids}
       return config, inputs_dict



@require_tf
class  A__       (	lowerCamelCase__	,       lowerCamelCase__	,       unittest.TestCase      ):
      _UpperCAmelCase     :Optional[int]						   =       (
          (
              TFDPRContextEncoder,
              TFDPRQuestionEncoder,
              TFDPRReader,
          )
          if is_tf_available()
          else ()
      )
      _UpperCAmelCase     :Optional[int]						   =       {'feature-extraction': TFDPRQuestionEncoder} if is_tf_available() else {}

      _UpperCAmelCase     :Any						   =       False
      _UpperCAmelCase     :str						   =       False
      _UpperCAmelCase     :Dict						   =       False
      _UpperCAmelCase     :Optional[int]						   =       False
      _UpperCAmelCase     :Dict						   =       False





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : Any	     =       TFDPRModelTester(self       )
       UpperCamelCase     : Dict	     =       ConfigTester(self	,							config_class=A_	,							hidden_size=37       )





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       self.config_tester.run_common_tests()





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : str	     =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_dpr_context_encoder(*A_       )





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : str	     =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_dpr_question_encoder(*A_       )





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : Union[str, Any]	     =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_dpr_reader(*A_       )





      @slow
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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

       for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCamelCase     : Union[str, Any]	     =       TFDPRContextEncoder.from_pretrained(A_       )
        self.assertIsNotNone(A_       )

       for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCamelCase     : Dict	     =       TFDPRQuestionEncoder.from_pretrained(A_       )
        self.assertIsNotNone(A_       )

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



@require_tf
class  A__       (	unittest.TestCase      ):





      @slow
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : int	     =       TFDPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base"       )

       UpperCamelCase     : Union[str, Any]	     =       tf.constant(
           [[101, 7592, 1010, 2003, 2026, 3899, 1_0140, 1029, 102]]       )  # [CLS] hello, is my dog cute? [SEP]
       UpperCamelCase     : List[Any]	     =       model(A_       )[0]  # embedding shape = (1, 768)
       # compare the actual values for a slice.
       UpperCamelCase     : str	     =       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       )       )







from __future__ import annotations

import unittest

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


if is_tf_available():
    import tensorflow as tf

    from transformers import AutoTokenizer, TFAutoModelForSeqaSeqLM



@require_tf
@require_sentencepiece
@require_tokenizers
class  A__       (	unittest.TestCase      ):





      @slow
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : List[str]	     =       TFAutoModelForSeqaSeqLM.from_pretrained("google/mt5-small"       )
       UpperCamelCase     : Tuple	     =       AutoTokenizer.from_pretrained("google/mt5-small"       )

       UpperCamelCase     : Dict	     =       tokenizer("Hello there"	,							return_tensors="tf"       ).input_ids
       UpperCamelCase     : int	     =       tokenizer("Hi I am"	,							return_tensors="tf"       ).input_ids

       UpperCamelCase     : Union[str, Any]	     =       model(A_	,							labels=A_       ).loss
       UpperCamelCase     : List[str]	     =       -tf.math.reduce_mean(A_       ).numpy()

       UpperCamelCase     : Union[str, Any]	     =       -21.22_81_68
       self.assertTrue(abs(mtf_score - EXPECTED_SCORE       ) < 2e-4       )
