





from __future__ import annotations

import unittest

from transformers import AutoTokenizer, MBartConfig, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow
from transformers.utils import cached_property

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 TFAutoModelForSeqaSeqLM, TFMBartForConditionalGeneration, TFMBartModel



@require_tf
class 					__snake_case	:


    '''simple docstring'''







    lowerCamelCase__    :		str				=	MBartConfig
    lowerCamelCase__    :		Any				=	{}
    lowerCamelCase__    :		List[Any]				=	"""gelu"""



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





        '''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_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



    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__											=			ids_tensor([self.batch_size, self.seq_length]   ,				self.vocab_size     )

        SCREAMING_SNAKE_CASE__											=			self.config_cls(
            vocab_size=self.vocab_size   ,				d_model=self.hidden_size   ,				encoder_layers=self.num_hidden_layers   ,				decoder_layers=self.num_hidden_layers   ,				encoder_attention_heads=self.num_attention_heads   ,				decoder_attention_heads=self.num_attention_heads   ,				encoder_ffn_dim=self.intermediate_size   ,				decoder_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_ids=[2]   ,				bos_token_id=self.bos_token_id   ,				pad_token_id=self.pad_token_id   ,				decoder_start_token_id=self.pad_token_id   ,				**self.config_updates   ,				)
        SCREAMING_SNAKE_CASE__											=			prepare_mbart_inputs_dict(A_   ,				A_   ,				A_     )
        return config, inputs_dict




    def  lowercase_					(       self   ,				A_   ,				A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			TFMBartModel(config=A_     ).get_decoder()
        SCREAMING_SNAKE_CASE__											=			inputs_dict['''input_ids''']

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

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

        SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__											=			outputs.to_tuple()
        SCREAMING_SNAKE_CASE__											=			past_key_values[1]





def 							__snake_case							(    lowerCAmelCase_ ,							lowerCAmelCase_ ,							lowerCAmelCase_ ,							lowerCAmelCase_=None ,							lowerCAmelCase_=None ,							lowerCAmelCase_=None ,							lowerCAmelCase_=None ,							lowerCAmelCase_=None ,							)				->						Dict:
    if attention_mask is None:
        SCREAMING_SNAKE_CASE__											=			tf.cast(tf.math.not_equal(lowerCAmelCase_ ,							config.pad_token_id	) ,							tf.inta	)
    if decoder_attention_mask is None:
        SCREAMING_SNAKE_CASE__											=			tf.concat(
            [
                tf.ones(decoder_input_ids[:, :1].shape ,							dtype=tf.inta	),
                tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:] ,							config.pad_token_id	) ,							tf.inta	),
            ] ,							axis=-1 ,							)
    if head_mask is None:
        SCREAMING_SNAKE_CASE__											=			tf.ones((config.encoder_layers, config.encoder_attention_heads)	)
    if decoder_head_mask is None:
        SCREAMING_SNAKE_CASE__											=			tf.ones((config.decoder_layers, config.decoder_attention_heads)	)
    if cross_attn_head_mask is None:
        SCREAMING_SNAKE_CASE__											=			tf.ones((config.decoder_layers, config.decoder_attention_heads)	)
    return {
        "input_ids": input_ids,
        "decoder_input_ids": decoder_input_ids,
        "attention_mask": attention_mask,
        "decoder_attention_mask": decoder_attention_mask,
        "head_mask": head_mask,
        "decoder_head_mask": decoder_head_mask,
        "cross_attn_head_mask": cross_attn_head_mask,
    }



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


    '''simple docstring'''







    lowerCamelCase__    :		Tuple				=	(TFMBartForConditionalGeneration, TFMBartModel) if is_tf_available() else ()
    lowerCamelCase__    :		List[str]				=	(TFMBartForConditionalGeneration,) if is_tf_available() else ()
    lowerCamelCase__    :		Any				=	(
        {
            """conversational""": TFMBartForConditionalGeneration,
            """feature-extraction""": TFMBartModel,
            """summarization""": TFMBartForConditionalGeneration,
            """text2text-generation""": TFMBartForConditionalGeneration,
            """translation""": TFMBartForConditionalGeneration,
        }
        if is_tf_available()
        else {}
    )
    lowerCamelCase__    :		List[str]				=	True
    lowerCamelCase__    :		Tuple				=	False
    lowerCamelCase__    :		Union[str, Any]				=	False



    def  lowercase_					(       self   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_     ):





        '''simple docstring'''

        if pipeline_test_casse_name != "FeatureExtractionPipelineTests":
            # Exception encountered when calling layer '...'
            return True

        return False



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			TFMBartModelTester(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_     )





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


    '''simple docstring'''







    lowerCamelCase__    :		Tuple				=	[
        """ UN Chief Says There Is No Military Solution in Syria""",
    ]
    lowerCamelCase__    :		Dict				=	[
        """Şeful ONU declară că nu există o soluţie militară în Siria""",
    ]
    lowerCamelCase__    :		Optional[int]				=	"""facebook/mbart-large-en-ro"""



    @cached_property
    def  lowercase_					(       self     ):





        '''simple docstring'''

        return AutoTokenizer.from_pretrained(self.model_name     )



    @cached_property
    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name     )
        return model



    def  lowercase_					(       self   ,				**A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.translate_src_text(**A_     )
        self.assertListEqual(self.expected_text   ,				A_     )



    def  lowercase_					(       self   ,				**A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.tokenizer(self.src_text   ,				**A_   ,				return_tensors='''tf'''     )
        SCREAMING_SNAKE_CASE__											=			self.model.generate(
            model_inputs.input_ids   ,				attention_mask=model_inputs.attention_mask   ,				num_beams=2     )
        SCREAMING_SNAKE_CASE__											=			self.tokenizer.batch_decode(A_   ,				skip_special_tokens=A_     )
        return generated_words




    @slow
    def  lowercase_					(       self     ):





        '''simple docstring'''

        self._assert_generated_batch_equal_expected()







from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Features, Value
from .base import TaskTemplate



@dataclass(frozen=__SCREAMING_SNAKE_CASE    )
class 					__snake_case	(  __SCREAMING_SNAKE_CASE    ):


    '''simple docstring'''







    lowerCamelCase__    :		str				=	field(default="""summarization"""			,     metadata={"""include_in_asdict_even_if_is_default""": True}    )
    lowerCamelCase__    :		ClassVar[Features]				=	Features({"""text""": Value("""string"""    )}    )
    lowerCamelCase__    :		ClassVar[Features]				=	Features({"""summary""": Value("""string"""    )}    )
    lowerCamelCase__    :		str				=	"text"
    lowerCamelCase__    :		str				=	"summary"



    @property
    def  lowercase_					(       self     ):





        '''simple docstring'''

        return {self.text_column: "text", self.summary_column: "summary"}
