





from __future__ import annotations

import unittest

import numpy as np

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

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


if is_tf_available():
  import tensorflow as tf

  from transformers import GPTaTokenizer, TFOPTForCausalLM, TFOPTModel


def 							__snake_case							(    lowerCAmelCase_ ,							lowerCAmelCase_ ,							lowerCAmelCase_=None ,							lowerCAmelCase_=None	)				->						Optional[int]:
    if attention_mask is None:
        SCREAMING_SNAKE_CASE__											=			tf.cast(tf.math.not_equal(lowerCAmelCase_ ,							config.pad_token_id	) ,							tf.inta	)
    return {"input_ids": input_ids, "attention_mask": attention_mask}



@require_tf
class 					__snake_case	:


    '''simple docstring'''







    lowerCamelCase__    :		Union[str, Any]				=	OPTConfig
    lowerCamelCase__    :		str				=	{}
    lowerCamelCase__    :		List[Any]				=	"""gelu"""



    def __init__(       self   ,				A_   ,				A_=13   ,				A_=7   ,				A_=True   ,				A_=False   ,				A_=99   ,				A_=16   ,				A_=2   ,				A_=4   ,				A_=4   ,				A_="gelu"   ,				A_=0.1   ,				A_=0.1   ,				A_=20   ,				A_=2   ,				A_=1   ,				A_=0   ,				A_=16   ,				A_=16   ,				):





        '''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_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__											=			eos_token_id
        SCREAMING_SNAKE_CASE__											=			pad_token_id
        SCREAMING_SNAKE_CASE__											=			bos_token_id
        SCREAMING_SNAKE_CASE__											=			embed_dim
        SCREAMING_SNAKE_CASE__											=			word_embed_proj_dim
        SCREAMING_SNAKE_CASE__											=			False



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			ids_tensor([self.batch_size, self.seq_length - 1]   ,				self.vocab_size     )
        SCREAMING_SNAKE_CASE__											=			tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size     )   ,				1     )
        SCREAMING_SNAKE_CASE__											=			tf.concat([input_ids, eos_tensor]   ,				axis=1     )

        SCREAMING_SNAKE_CASE__											=			self.config_cls(
            vocab_size=self.vocab_size   ,				hidden_size=self.hidden_size   ,				num_hidden_layers=self.num_hidden_layers   ,				num_attention_heads=self.num_attention_heads   ,				ffn_dim=self.intermediate_size   ,				dropout=self.hidden_dropout_prob   ,				attention_dropout=self.attention_probs_dropout_prob   ,				max_position_embeddings=self.max_position_embeddings   ,				eos_token_id=self.eos_token_id   ,				bos_token_id=self.bos_token_id   ,				pad_token_id=self.pad_token_id   ,				embed_dim=self.embed_dim   ,				word_embed_proj_dim=self.word_embed_proj_dim   ,				is_encoder_decoder=A_   ,				**self.config_updates   ,				)
        SCREAMING_SNAKE_CASE__											=			prepare_opt_inputs_dict(A_   ,				A_     )
        return config, inputs_dict




    def  lowercase_					(       self   ,				A_   ,				A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			TFOPTModel(config=A_     )
        SCREAMING_SNAKE_CASE__											=			inputs_dict['''input_ids''']

        SCREAMING_SNAKE_CASE__											=			input_ids[:1, :]
        SCREAMING_SNAKE_CASE__											=			inputs_dict['''attention_mask'''][:1, :]
        SCREAMING_SNAKE_CASE__											=			1

        # first forward pass
        SCREAMING_SNAKE_CASE__											=			model(A_   ,				attention_mask=A_   ,				use_cache=A_     )

        SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__											=			outputs.to_tuple()

        # create hypothetical next token and extent to next_input_ids
        SCREAMING_SNAKE_CASE__											=			ids_tensor((self.batch_size, 3)   ,				config.vocab_size     )
        SCREAMING_SNAKE_CASE__											=			tf.cast(ids_tensor((self.batch_size, 3)   ,				2     )   ,				tf.inta     )

        # append to next input_ids and
        SCREAMING_SNAKE_CASE__											=			tf.concat([input_ids, next_tokens]   ,				axis=-1     )
        SCREAMING_SNAKE_CASE__											=			tf.concat([attention_mask, next_attn_mask]   ,				axis=-1     )

        SCREAMING_SNAKE_CASE__											=			model(A_   ,				attention_mask=A_     )[0]
        SCREAMING_SNAKE_CASE__											=			model(A_   ,				attention_mask=A_   ,				past_key_values=A_     )[0]

        self.parent.assertEqual(next_tokens.shape[1]   ,				output_from_past.shape[1]     )

        # select random slice
        SCREAMING_SNAKE_CASE__											=			int(ids_tensor((1,)   ,				output_from_past.shape[-1]     )     )
        SCREAMING_SNAKE_CASE__											=			output_from_no_past[:, -3:, random_slice_idx]
        SCREAMING_SNAKE_CASE__											=			output_from_past[:, :, random_slice_idx]

        # test that outputs are equal for slice
        tf.debugging.assert_near(A_   ,				A_   ,				rtol=1E-3     )





@require_tf
class 					__snake_case	(  __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     unittest.TestCase    ):


    '''simple docstring'''







    lowerCamelCase__    :		Dict				=	(TFOPTModel, TFOPTForCausalLM) if is_tf_available() else ()
    lowerCamelCase__    :		List[str]				=	(TFOPTForCausalLM,) if is_tf_available() else ()
    lowerCamelCase__    :		Optional[int]				=	(
        {"""feature-extraction""": TFOPTModel, """text-generation""": TFOPTForCausalLM} if is_tf_available() else {}
    )
    lowerCamelCase__    :		int				=	False
    lowerCamelCase__    :		str				=	False
    lowerCamelCase__    :		str				=	False
    lowerCamelCase__    :		Tuple				=	1_0



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			TFOPTModelTester(self     )
        SCREAMING_SNAKE_CASE__											=			ConfigTester(self   ,				config_class=A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        self.config_tester.run_common_tests()



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.check_decoder_model_past_large_inputs(*A_     )




    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs_for_common()

        def _get_word_embedding_weight(A_   ,				A_     ):
            if hasattr(A_   ,				'''weight'''     ):
                return embedding_layer.weight
            else:
                # Here we build the word embeddings weights if not exists.
                # And then we retry to get the attribute once built.
                model.build()
                if hasattr(A_   ,				'''weight'''     ):
                    return embedding_layer.weight
                else:
                    return None

        for model_class in self.all_model_classes:
            for size in [config.vocab_size - 10, config.vocab_size + 10]:
                # build the embeddings
                SCREAMING_SNAKE_CASE__											=			model_class(config=A_     )
                SCREAMING_SNAKE_CASE__											=			_get_word_embedding_weight(A_   ,				model.get_input_embeddings()     )
                SCREAMING_SNAKE_CASE__											=			_get_word_embedding_weight(A_   ,				model.get_output_embeddings()     )

                # reshape the embeddings
                model.resize_token_embeddings(A_     )
                SCREAMING_SNAKE_CASE__											=			_get_word_embedding_weight(A_   ,				model.get_input_embeddings()     )
                SCREAMING_SNAKE_CASE__											=			_get_word_embedding_weight(A_   ,				model.get_output_embeddings()     )

                # check that the resized embeddings size matches the desired size.
                SCREAMING_SNAKE_CASE__											=			size if size is not None else config.vocab_size

                self.assertEqual(new_input_embeddings.shape[0]   ,				A_     )

                # check that weights remain the same after resizing
                SCREAMING_SNAKE_CASE__											=			True
                for pa, pa in zip(old_input_embeddings.value()   ,				new_input_embeddings.value()     ):
                    if tf.math.reduce_sum(tf.math.abs(pa - pa     )     ) > 0:
                        SCREAMING_SNAKE_CASE__											=			False
                self.assertTrue(A_     )

                if old_output_embeddings is not None and new_output_embeddings is not None:
                    self.assertEqual(new_output_embeddings.shape[0]   ,				A_     )

                    SCREAMING_SNAKE_CASE__											=			True
                    for pa, pa in zip(old_output_embeddings.value()   ,				new_output_embeddings.value()     ):
                        if tf.math.reduce_sum(tf.math.abs(pa - pa     )     ) > 0:
                            SCREAMING_SNAKE_CASE__											=			False
                    self.assertTrue(A_     )





def 							__snake_case							(    lowerCAmelCase_	)				->						List[str]:
    return tf.constant(lowerCAmelCase_ ,							dtype=tf.intaa	)



@require_tf
class 					__snake_case	(  unittest.TestCase    ):


    '''simple docstring'''







    lowerCamelCase__    :		Optional[Any]				=	9_9



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			tf.ones((4, 1)   ,				dtype=tf.intaa     ) * 2
        SCREAMING_SNAKE_CASE__											=			tf.concat([ids_tensor((4, 6)   ,				self.vocab_size - 3     ) + 3, eos_column_vector]   ,				axis=1     )
        SCREAMING_SNAKE_CASE__											=			input_ids.shape[0]
        SCREAMING_SNAKE_CASE__											=			OPTConfig(
            vocab_size=self.vocab_size   ,				hidden_size=24   ,				num_hidden_layers=2   ,				num_attention_heads=2   ,				ffn_dim=32   ,				max_position_embeddings=48   ,				eos_token_id=2   ,				pad_token_id=1   ,				bos_token_id=0   ,				)
        return config, input_ids, batch_size





@require_sentencepiece
@require_tf
class 					__snake_case	(  unittest.TestCase    ):


    '''simple docstring'''







    @slow
    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			TFOPTModel.from_pretrained('''facebook/opt-350m'''     )
        SCREAMING_SNAKE_CASE__											=			_long_tensor([[0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 4_60_78, 15_88, 2]]     )
        SCREAMING_SNAKE_CASE__											=			tf.not_equal(A_   ,				model.config.pad_token_id     )
        with tf.GradientTape():
            SCREAMING_SNAKE_CASE__											=			model(input_ids=A_   ,				attention_mask=A_     ).last_hidden_state
        SCREAMING_SNAKE_CASE__											=			(1, 11, 5_12)
        self.assertEqual(output.shape   ,				A_     )
        SCREAMING_SNAKE_CASE__											=			tf.constant(
            [[-0.2873, -1.9218, -0.3033], [-1.2710, -0.1338, -0.1902], [0.4095, 0.1214, -1.3121]]     )
        self.assertTrue(np.allclose(output[:, :3, :3]   ,				A_   ,				atol=4E-3     )     )

        SCREAMING_SNAKE_CASE__											=			tf.function(A_   ,				jit_compile=A_     )
        SCREAMING_SNAKE_CASE__											=			xla_generate(A_   ,				A_     )[0]
        self.assertTrue(np.allclose(output[:, :3, :3]   ,				A_   ,				atol=4E-2     )     )





@require_tf
@slow
class 					__snake_case	(  unittest.TestCase    ):


    '''simple docstring'''







    def  lowercase_					(       self     ):





        '''simple docstring'''

        super().setUp()
        SCREAMING_SNAKE_CASE__											=			'''facebook/opt-350m'''




    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			TFOPTForCausalLM.from_pretrained(self.path_model     )
        SCREAMING_SNAKE_CASE__											=			GPTaTokenizer.from_pretrained(self.path_model     )

        SCREAMING_SNAKE_CASE__											=			[
            '''Today is a beautiful day and I want to''',
            '''In the city of''',
            '''Paris is the capital of France and''',
            '''Computers and mobile phones have taken''',
        ]
        # verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
        SCREAMING_SNAKE_CASE__											=			tokenizer(A_   ,				return_tensors='''tf'''   ,				padding=A_   ,				add_special_tokens=A_     )
        SCREAMING_SNAKE_CASE__											=			tf.math.reduce_mean(model(inputs.input_ids   ,				attention_mask=inputs.attention_mask     )[0]   ,				axis=-1     )
        SCREAMING_SNAKE_CASE__											=			tf.constant(
            [
                [1.3851, -13.8923, -10.5229, -10.7533, -0.2309, -10.2384, -0.5365, -9.0947, -5.1670],
                [-4.7073, -10.6276, -3.9415, -21.5242, -0.2822, -0.2822, -0.2822, -0.2822, -0.2822],
                [0.6247, -3.4229, -8.9179, -1.4297, -14.1650, 1.4146, -9.0218, -0.2703, -0.2703],
                [6.4783, -1.9913, -10.7926, -2.3336, 1.5092, -0.9974, -6.8213, 1.3477, 1.3477],
            ]     )
        self.assertTrue(np.allclose(A_   ,				A_   ,				atol=1E-4     )     )

        SCREAMING_SNAKE_CASE__											=			tf.function(A_   ,				jit_compile=A_     )
        SCREAMING_SNAKE_CASE__											=			tf.math.reduce_mean(xla_generate(inputs.input_ids   ,				attention_mask=inputs.attention_mask     )[0]   ,				axis=-1     )
        self.assertTrue(np.allclose(A_   ,				A_   ,				atol=1E-4     )     )





@require_tf
@slow
class 					__snake_case	(  unittest.TestCase    ):


    '''simple docstring'''







    @property
    def  lowercase_					(       self     ):





        '''simple docstring'''

        return [
            "Today is a beautiful day and I want",
            "In the city of",
            "Paris is the capital of France and",
            "Computers and mobile phones have taken",
        ]



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			'''facebook/opt-125m'''

        SCREAMING_SNAKE_CASE__											=			[
            '''Today is a beautiful day and I want to''',
            '''In the city of New York, the city''',
            '''Paris is the capital of France and the capital''',
            '''Computers and mobile phones have taken over the''',
        ]

        SCREAMING_SNAKE_CASE__											=			[]
        SCREAMING_SNAKE_CASE__											=			GPTaTokenizer.from_pretrained(A_     )
        SCREAMING_SNAKE_CASE__											=			TFOPTForCausalLM.from_pretrained(A_     )

        for prompt in self.prompts:
            SCREAMING_SNAKE_CASE__											=			tokenizer(A_   ,				return_tensors='''tf'''     ).input_ids

            SCREAMING_SNAKE_CASE__											=			model.generate(A_   ,				max_length=10     )

            SCREAMING_SNAKE_CASE__											=			tokenizer.batch_decode(A_   ,				skip_special_tokens=A_     )
            predicted_outputs += generated_string

        self.assertListEqual(A_   ,				A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			'''facebook/opt-350m'''

        SCREAMING_SNAKE_CASE__											=			GPTaTokenizer.from_pretrained(A_     )
        SCREAMING_SNAKE_CASE__											=			TFOPTForCausalLM.from_pretrained(A_     )

        SCREAMING_SNAKE_CASE__											=			'''left'''

        # use different length sentences to test batching
        SCREAMING_SNAKE_CASE__											=			[
            '''Hello, my dog is a little''',
            '''Today, I''',
        ]

        SCREAMING_SNAKE_CASE__											=			tokenizer(A_   ,				return_tensors='''tf'''   ,				padding=A_     )
        SCREAMING_SNAKE_CASE__											=			inputs['''input_ids''']

        SCREAMING_SNAKE_CASE__											=			model.generate(input_ids=A_   ,				attention_mask=inputs['''attention_mask''']     )

        SCREAMING_SNAKE_CASE__											=			tokenizer(sentences[0]   ,				return_tensors='''tf'''     ).input_ids
        SCREAMING_SNAKE_CASE__											=			model.generate(input_ids=A_     )

        SCREAMING_SNAKE_CASE__											=			inputs_non_padded.shape[-1] - tf.math.reduce_sum(
            tf.cast(inputs['''attention_mask'''][-1]   ,				tf.intaa     )     )
        SCREAMING_SNAKE_CASE__											=			tokenizer(sentences[1]   ,				return_tensors='''tf'''     ).input_ids
        SCREAMING_SNAKE_CASE__											=			model.generate(input_ids=A_   ,				max_length=model.config.max_length - num_paddings     )

        SCREAMING_SNAKE_CASE__											=			tokenizer.batch_decode(A_   ,				skip_special_tokens=A_     )
        SCREAMING_SNAKE_CASE__											=			tokenizer.decode(output_non_padded[0]   ,				skip_special_tokens=A_     )
        SCREAMING_SNAKE_CASE__											=			tokenizer.decode(output_padded[0]   ,				skip_special_tokens=A_     )

        SCREAMING_SNAKE_CASE__											=			[
            '''Hello, my dog is a little bit of a dork.\nI\'m a little bit''',
            '''Today, I was in the middle of a conversation with a friend about the''',
        ]
        self.assertListEqual(A_   ,				A_     )
        self.assertListEqual(A_   ,				[non_padded_sentence, padded_sentence]     )




    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			'''facebook/opt-350m'''

        SCREAMING_SNAKE_CASE__											=			[
            '''Today is a beautiful day and I want to''',
            '''In the city of San Francisco, the city''',
            '''Paris is the capital of France and the capital''',
            '''Computers and mobile phones have taken over the''',
        ]

        SCREAMING_SNAKE_CASE__											=			[]
        SCREAMING_SNAKE_CASE__											=			GPTaTokenizer.from_pretrained(A_     )
        SCREAMING_SNAKE_CASE__											=			TFOPTForCausalLM.from_pretrained(A_     )

        for prompt in self.prompts:
            SCREAMING_SNAKE_CASE__											=			tokenizer(A_   ,				return_tensors='''tf'''     ).input_ids

            SCREAMING_SNAKE_CASE__											=			model.generate(A_   ,				max_length=10     )

            SCREAMING_SNAKE_CASE__											=			tokenizer.batch_decode(A_   ,				skip_special_tokens=A_     )
            predicted_outputs += generated_string

        self.assertListEqual(A_   ,				A_     )






'''simple docstring'''
from ...configuration_utils import PretrainedConfig
from ...utils import logging


_SCREAMING_SNAKE_CASE			=						logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE			=						{
    "facebook/dpr-ctx_encoder-single-nq-base": (
        "https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/config.json"
    ),
    "facebook/dpr-question_encoder-single-nq-base": (
        "https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/config.json"
    ),
    "facebook/dpr-reader-single-nq-base": (
        "https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/config.json"
    ),
    "facebook/dpr-ctx_encoder-multiset-base": (
        "https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/config.json"
    ),
    "facebook/dpr-question_encoder-multiset-base": (
        "https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/config.json"
    ),
    "facebook/dpr-reader-multiset-base": (
        "https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/config.json"
    ),
}
class   __UpperCAmelCase				(     lowerCAmelCase				):



					'''simple docstring'''





					_UpperCamelCase							=		"""dpr"""



					def __init__(	self     :						Optional[Any]      ,       _lowercase     :						Dict=30_522      ,       _lowercase     :						Tuple=768      ,       _lowercase     :						int=12      ,       _lowercase     :						List[str]=12      ,       _lowercase     :						Union[str, Any]=3_072      ,       _lowercase     :						Dict="gelu"      ,       _lowercase     :						Dict=0.1      ,       _lowercase     :						Any=0.1      ,       _lowercase     :						Union[str, Any]=512      ,       _lowercase     :						Optional[int]=2      ,       _lowercase     :						Union[str, Any]=0.02      ,       _lowercase     :						List[Any]=1E-12      ,       _lowercase     :						List[str]=0      ,       _lowercase     :						Dict="absolute"      ,       _lowercase     :						int = 0      ,       **_lowercase     :						int      ,       )    ->							int:
								super().__init__(pad_token_id=_lowercase      ,       **_lowercase)

								A_									= vocab_size
								A_									= hidden_size
								A_									= num_hidden_layers
								A_									= num_attention_heads
								A_									= hidden_act
								A_									= intermediate_size
								A_									= hidden_dropout_prob
								A_									= attention_probs_dropout_prob
								A_									= max_position_embeddings
								A_									= type_vocab_size
								A_									= initializer_range
								A_									= layer_norm_eps
								A_									= projection_dim
								A_									= position_embedding_type




