



from typing import Callable, Optional, Union

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


UpperCamelCase__					    =    logging.get_logger(__name__)

UpperCamelCase__					    =    {
    """microsoft/xprophetnet-large-wiki100-cased""": (
        """https://huggingface.co/microsoft/xprophetnet-large-wiki100-cased/resolve/main/config.json"""
    ),
}
class     a__					(      snake_case__   ):
						_a							:					Any      =					"""xlm-prophetnet"""
						_a							:					Any      =					["""past_key_values"""]
						_a							:					List[Any]      =					{
						    """num_attention_heads""": """num_encoder_attention_heads""",
						}






						def __init__(     self	,    _A = 0.1	,    _A = "gelu"	,    _A = 3_0_5_2_2	,    _A = 1_0_2_4	,    _A = 4_0_9_6	,    _A = 1_2	,    _A = 1_6	,    _A = 4_0_9_6	,    _A = 1_2	,    _A = 1_6	,    _A = 0.1	,    _A = 0.1	,    _A = 5_1_2	,    _A = 0.02	,    _A = True	,    _A = True	,    _A = 0	,    _A = 2	,    _A = 3_2	,    _A = 1_2_8	,    _A = False	,    _A = 0.0	,    _A = True	,    _A = 0	,    _A = 1	,    _A = 2	,    **_A	,    ):




											"""simple docstring"""



											__lowerCAmelCase            = vocab_size
											__lowerCAmelCase            = hidden_size
											__lowerCAmelCase            = encoder_ffn_dim
											__lowerCAmelCase            = num_encoder_layers
											__lowerCAmelCase            = num_encoder_attention_heads
											__lowerCAmelCase            = decoder_ffn_dim
											__lowerCAmelCase            = num_decoder_layers
											__lowerCAmelCase            = num_decoder_attention_heads
											__lowerCAmelCase            = max_position_embeddings
											__lowerCAmelCase            = init_std  # Normal(0, this parameter)
											__lowerCAmelCase            = activation_function

											# parameters for xlmprophetnet
											__lowerCAmelCase            = ngram
											__lowerCAmelCase            = num_buckets
											__lowerCAmelCase            = relative_max_distance
											__lowerCAmelCase            = disable_ngram_loss
											__lowerCAmelCase            = eps

											# 3 Types of Dropout
											__lowerCAmelCase            = attention_dropout
											__lowerCAmelCase            = activation_dropout
											__lowerCAmelCase            = dropout

											__lowerCAmelCase            = use_cache

											super().__init__(
											    pad_token_id=_A	,    bos_token_id=_A	,    eos_token_id=_A	,    is_encoder_decoder=_A	,    add_cross_attention=_A	,    decoder_start_token_id=_A	,    **_A	,    )






						@property
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											return self.num_encoder_layers + self.num_decoder_layers






						@num_hidden_layers.setter
						def        __SCREAMING_SNAKE_CASE(     self	,    _A    ):




											"""simple docstring"""



											raise NotImplementedError(
											    "This model does not support the setting of `num_hidden_layers`. Please set `num_encoder_layers` and"
											    " `num_decoder_layers`."    )










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     a__					:
						_a							:					List[str]      =					MBartConfig
						_a							:					Union[str, Any]      =					{}
						_a							:					Tuple      =					"""gelu"""






						def __init__(     self	,    _A	,    _A=1_3	,    _A=7	,    _A=True	,    _A=False	,    _A=9_9	,    _A=3_2	,    _A=2	,    _A=4	,    _A=3_7	,    _A=0.1	,    _A=0.1	,    _A=2_0	,    _A=2	,    _A=1	,    _A=0	,    ):




											"""simple docstring"""



											__lowerCAmelCase            = parent
											__lowerCAmelCase            = batch_size
											__lowerCAmelCase            = seq_length
											__lowerCAmelCase            = is_training
											__lowerCAmelCase            = use_labels
											__lowerCAmelCase            = vocab_size
											__lowerCAmelCase            = hidden_size
											__lowerCAmelCase            = num_hidden_layers
											__lowerCAmelCase            = num_attention_heads
											__lowerCAmelCase            = intermediate_size
											__lowerCAmelCase            = hidden_dropout_prob
											__lowerCAmelCase            = attention_probs_dropout_prob
											__lowerCAmelCase            = max_position_embeddings
											__lowerCAmelCase            = eos_token_id
											__lowerCAmelCase            = pad_token_id
											__lowerCAmelCase            = bos_token_id






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



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

											__lowerCAmelCase            = ids_tensor([self.batch_size, self.seq_length]	,    self.vocab_size    )

											__lowerCAmelCase            = 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	,    )
											__lowerCAmelCase            = prepare_mbart_inputs_dict(_A	,    _A	,    _A    )
											return config, inputs_dict






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFMBartModel(config=_A    ).get_decoder()
											__lowerCAmelCase            = inputs_dict["input_ids"]

											__lowerCAmelCase            = input_ids[:1, :]
											__lowerCAmelCase            = inputs_dict["attention_mask"][:1, :]
											__lowerCAmelCase            = inputs_dict["head_mask"]
											__lowerCAmelCase            = 1

											# first forward pass
											__lowerCAmelCase            = model(_A	,    attention_mask=_A	,    head_mask=_A	,    use_cache=_A    )

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







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Optional[int]   ,   SCREAMING_SNAKE_CASE_							:       List[str]   ,   SCREAMING_SNAKE_CASE_							:       Tuple   ,   SCREAMING_SNAKE_CASE_							:       Any=None   ,   SCREAMING_SNAKE_CASE_							:       int=None   ,   SCREAMING_SNAKE_CASE_							:       List[Any]=None   ,   SCREAMING_SNAKE_CASE_							:       List[Any]=None   ,   SCREAMING_SNAKE_CASE_							:       Optional[Any]=None   ,   ):
					if attention_mask is None:
										__lowerCAmelCase            = tf.cast(tf.math.not_equal(SCREAMING_SNAKE_CASE_   ,   config.pad_token_id  )   ,   tf.inta  )
					if decoder_attention_mask is None:
										__lowerCAmelCase            = 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:
										__lowerCAmelCase            = tf.ones((config.encoder_layers, config.encoder_attention_heads)  )
					if decoder_head_mask is None:
										__lowerCAmelCase            = tf.ones((config.decoder_layers, config.decoder_attention_heads)  )
					if cross_attn_head_mask is None:
										__lowerCAmelCase            = 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     a__					(      snake_case__				,       snake_case__				,       unittest.TestCase   ):
						_a							:					Optional[Any]      =					(TFMBartForConditionalGeneration, TFMBartModel) if is_tf_available() else ()
						_a							:					Any      =					(TFMBartForConditionalGeneration,) if is_tf_available() else ()
						_a							:					Tuple      =					(
						    {
						        """conversational""": TFMBartForConditionalGeneration,
						        """feature-extraction""": TFMBartModel,
						        """summarization""": TFMBartForConditionalGeneration,
						        """text2text-generation""": TFMBartForConditionalGeneration,
						        """translation""": TFMBartForConditionalGeneration,
						    }
						    if is_tf_available()
						    else {}
						)
						_a							:					Optional[int]      =					True
						_a							:					Optional[int]      =					False
						_a							:					List[Any]      =					False






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A	,    _A	,    _A    ):




											"""simple docstring"""



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

											return False






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = TFMBartModelTester(self    )
											__lowerCAmelCase            = ConfigTester(self	,    config_class=_A    )






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											self.config_tester.run_common_tests()






						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = 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     a__					(      unittest.TestCase   ):
						_a							:					Any      =					[
						    """ UN Chief Says There Is No Military Solution in Syria""",
						]
						_a							:					Optional[Any]      =					[
						    """Şeful ONU declară că nu există o soluţie militară în Siria""",
						]
						_a							:					Any      =					"""facebook/mbart-large-en-ro"""






						@cached_property
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											return AutoTokenizer.from_pretrained(self.model_name    )






						@cached_property
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



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






						def        __SCREAMING_SNAKE_CASE(     self	,    **_A    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.translate_src_text(**_A    )
											self.assertListEqual(self.expected_text	,    _A    )






						def        __SCREAMING_SNAKE_CASE(     self	,    **_A    ):




											"""simple docstring"""



											__lowerCAmelCase            = self.tokenizer(self.src_text	,    **_A	,    return_tensors="tf"    )
											__lowerCAmelCase            = self.model.generate(
											    model_inputs.input_ids	,    attention_mask=model_inputs.attention_mask	,    num_beams=2    )
											__lowerCAmelCase            = self.tokenizer.batch_decode(_A	,    skip_special_tokens=_A    )
											return generated_words






						@slow
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											self._assert_generated_batch_equal_expected()




