




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			:
   UpperCAmelCase__  :    List[Any]          =  MBartConfig
   UpperCAmelCase__  :    Optional[Any]          =  {}
   UpperCAmelCase__  :    Optional[Any]          =  '''gelu'''







   def __init__(      self		:						Dict					,						_snake_case		:						Any					,						_snake_case		:						int=13					,						_snake_case		:						Union[str, Any]=7					,						_snake_case		:						Tuple=True					,						_snake_case		:						Tuple=False					,						_snake_case		:						Dict=99					,						_snake_case		:						str=32					,						_snake_case		:						List[str]=2					,						_snake_case		:						int=4					,						_snake_case		:						List[str]=37					,						_snake_case		:						int=0.1					,						_snake_case		:						Dict=0.1					,						_snake_case		:						Dict=20					,						_snake_case		:						Tuple=2					,						_snake_case		:						Tuple=1					,						_snake_case		:						List[Any]=0					,						):

        """simple docstring"""


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







   def lowerCamelCase						(      self		:						str):

        """simple docstring"""


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

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

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







   def lowerCamelCase						(      self		:						Tuple					,						_snake_case		:						List[str]					,						_snake_case		:						Union[str, Any]):

        """simple docstring"""


        UpperCAmelCase_          =							TFMBartModel(config=_snake_case).get_decoder()
        UpperCAmelCase_          =							inputs_dict['''input_ids''']

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

        # first forward pass
        UpperCAmelCase_          =							model(_snake_case					,						attention_mask=_snake_case					,						head_mask=_snake_case					,						use_cache=_snake_case)

        UpperCAmelCase_          =							outputs.to_tuple()
        UpperCAmelCase_          =							past_key_values[1]







def    A						(__A   :    Optional[int]   ,				__A   :    Union[str, Any]   ,				__A   :    List[Any]   ,				__A   :    Union[str, Any]=None   ,				__A   :    List[Any]=None   ,				__A   :    Any=None   ,				__A   :    Union[str, Any]=None   ,				__A   :    str=None   ,				)   ->      Union[str, Any]:







     """simple docstring"""


     if attention_mask is None:
          UpperCAmelCase_          =							tf.cast(tf.math.not_equal(UpperCamelCase__   ,				config.pad_token_id       )   ,				tf.inta       )
     if decoder_attention_mask is None:
          UpperCAmelCase_          =							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:
          UpperCAmelCase_          =							tf.ones((config.encoder_layers, config.encoder_attention_heads)       )
     if decoder_head_mask is None:
          UpperCAmelCase_          =							tf.ones((config.decoder_layers, config.decoder_attention_heads)       )
     if cross_attn_head_mask is None:
          UpperCAmelCase_          =							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			(			a ,  a ,  unittest.TestCase		):
   UpperCAmelCase__  :    Tuple          =  (TFMBartForConditionalGeneration, TFMBartModel) if is_tf_available() else ()
   UpperCAmelCase__  :    Tuple          =  (TFMBartForConditionalGeneration,) if is_tf_available() else ()
   UpperCAmelCase__  :    Optional[Any]          =  (
       {
           '''conversational''': TFMBartForConditionalGeneration,
           '''feature-extraction''': TFMBartModel,
           '''summarization''': TFMBartForConditionalGeneration,
           '''text2text-generation''': TFMBartForConditionalGeneration,
           '''translation''': TFMBartForConditionalGeneration,
       }
       if is_tf_available()
       else {}
   )
   UpperCAmelCase__  :    Union[str, Any]          =  True
   UpperCAmelCase__  :    Union[str, Any]          =  False
   UpperCAmelCase__  :    Any          =  False







   def lowerCamelCase						(      self		:						Union[str, Any]					,						_snake_case		:						Optional[int]					,						_snake_case		:						List[str]					,						_snake_case		:						List[str]					,						_snake_case		:						List[Any]					,						_snake_case		:						int):

        """simple docstring"""


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

        return False







   def lowerCamelCase						(      self		:						Dict):

        """simple docstring"""


        UpperCAmelCase_          =							TFMBartModelTester(self)
        UpperCAmelCase_          =							ConfigTester(self					,						config_class=_snake_case)







   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        self.config_tester.run_common_tests()







   def lowerCamelCase						(      self		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.check_decoder_model_past_large_inputs(*_snake_case)







@require_sentencepiece
@require_tokenizers
@require_tf
class __snake_case			(			unittest.TestCase		):
   UpperCAmelCase__  :    List[str]          =  [
       ''' UN Chief Says There Is No Military Solution in Syria''',
   ]
   UpperCAmelCase__  :    str          =  [
       '''Şeful ONU declară că nu există o soluţie militară în Siria''',
   ]
   UpperCAmelCase__  :    Union[str, Any]          =  '''facebook/mbart-large-en-ro'''







   @cached_property
   def lowerCamelCase						(      self		:						Union[str, Any]):

        """simple docstring"""


        return AutoTokenizer.from_pretrained(self.model_name)







   @cached_property
   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


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







   def lowerCamelCase						(      self		:						List[str]					,						**_snake_case		:						List[Any]):

        """simple docstring"""


        UpperCAmelCase_          =							self.translate_src_text(**_snake_case)
        self.assertListEqual(self.expected_text					,						_snake_case)







   def lowerCamelCase						(      self		:						Union[str, Any]					,						**_snake_case		:						str):

        """simple docstring"""


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







   @slow
   def lowerCamelCase						(      self		:						Optional[int]):

        """simple docstring"""


        self._assert_generated_batch_equal_expected()











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


_lowerCAmelCase    :List[str]     =					logging.get_logger(__name__)

_lowerCAmelCase    :Any     =					{
    'tiiuae/falcon-40b': 'https://huggingface.co/tiiuae/falcon-40b/resolve/main/config.json',
    'tiiuae/falcon-7b': 'https://huggingface.co/tiiuae/falcon-7b/resolve/main/config.json',
}
class 							_UpperCAmelCase							( a      ):





     '''simple docstring'''



     a__									='''falcon'''
     a__									=['''past_key_values''']





     def __init__(					self       ,	A=6_5_0_2_4       ,	A=4_5_4_4       ,	A=3_2       ,	A=7_1       ,	A=1E-5       ,	A=0.02       ,	A=True       ,	A=0.0       ,	A=0.0       ,	A=None       ,	A=False       ,	A=False       ,	A=True       ,	A=True       ,	A=False       ,	A=1_1       ,	A=1_1       ,	**A       ,	)							->					Any:
         _UpperCAmelCase							:  int											=    vocab_size
         # Backward compatibility with n_embed kwarg
         _UpperCAmelCase							:  Optional[Any]											=    kwargs.pop('''n_embed'''       ,	A		)
         _UpperCAmelCase							:  int											=    hidden_size if n_embed is None else n_embed
         _UpperCAmelCase							:  List[str]											=    num_hidden_layers
         _UpperCAmelCase							:  Tuple											=    num_attention_heads
         _UpperCAmelCase							:  Optional[int]											=    layer_norm_epsilon
         _UpperCAmelCase							:  Tuple											=    initializer_range
         _UpperCAmelCase							:  Optional[int]											=    use_cache
         _UpperCAmelCase							:  Any											=    hidden_dropout
         _UpperCAmelCase							:  Dict											=    attention_dropout

         _UpperCAmelCase							:  Any											=    bos_token_id
         _UpperCAmelCase							:  List[Any]											=    eos_token_id
         _UpperCAmelCase							:  Tuple											=    num_attention_heads if num_kv_heads is None else num_kv_heads
         _UpperCAmelCase							:  Dict											=    alibi
         _UpperCAmelCase							:  Optional[int]											=    new_decoder_architecture
         _UpperCAmelCase							:  str											=    multi_query  # Ignored when new_decoder_architecture is True
         _UpperCAmelCase							:  Optional[int]											=    parallel_attn
         _UpperCAmelCase							:  Optional[int]											=    bias

         super().__init__(bos_token_id=A       ,	eos_token_id=A       ,	**A		)





     @property
     def  __lowerCAmelCase	(					self		)							->					List[str]:
         return self.hidden_size // self.num_attention_heads





     @property
     def  __lowerCAmelCase	(					self		)							->					List[Any]:
         return not self.alibi


