

"""simple docstring"""





import tensorflow as tf

from ...tf_utils import shape_list







class 				SCREAMING_SNAKE_CASE					(		tf.keras.layers.Layer       ):






					"""simple docstring"""



					def __init__(       self			:				List[Any]   ,lowercase_			:				Tuple   ,lowercase_			:				List[Any]   ,lowercase_			:				List[str]   ,lowercase_			:				Optional[int]   ,lowercase_			:				str=1   ,lowercase_			:				Union[str, Any]=False   ,**lowercase_			:				Union[str, Any]		):
						super().__init__(**lowercase_		)

						lowerCAmelCase__	:						Any  				=	vocab_size
						lowerCAmelCase__	:						Tuple  				=	d_embed
						lowerCAmelCase__	:						Union[str, Any]  				=	d_proj

						lowerCAmelCase__	:						Union[str, Any]  				=	cutoffs + [vocab_size]
						lowerCAmelCase__	:						Optional[Any]  				=	[0] + self.cutoffs
						lowerCAmelCase__	:						int  				=	div_val

						lowerCAmelCase__	:						Optional[Any]  				=	self.cutoffs[0]
						lowerCAmelCase__	:						Tuple  				=	len(self.cutoffs		) - 1
						lowerCAmelCase__	:						Any  				=	self.shortlist_size + self.n_clusters
						lowerCAmelCase__	:						str  				=	keep_order

						lowerCAmelCase__	:						List[str]  				=	[]
						lowerCAmelCase__	:						List[str]  				=	[]







					def        __lowerCAmelCase (       self			:				Optional[int]   ,lowercase_			:				str		):
						if self.n_clusters > 0:
							lowerCAmelCase__	:						Optional[Any]  				=	self.add_weight(
							    shape=(self.n_clusters, self.d_embed)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name='''cluster_weight'''		)
							lowerCAmelCase__	:						List[Any]  				=	self.add_weight(
							    shape=(self.n_clusters,)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name='''cluster_bias'''		)

						if self.div_val == 1:
							for i in range(len(self.cutoffs		)		):
								if self.d_proj != self.d_embed:
									lowerCAmelCase__	:						Optional[Any]  				=	self.add_weight(
									    shape=(self.d_embed, self.d_proj)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name=F'out_projs_._{i}'   ,)
									self.out_projs.append(lowercase_		)
								else:
									self.out_projs.append(lowercase_		)
								lowerCAmelCase__	:						Any  				=	self.add_weight(
								    shape=(self.vocab_size, self.d_embed)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name=F'out_layers_._{i}_._weight'   ,)
								lowerCAmelCase__	:						Union[str, Any]  				=	self.add_weight(
								    shape=(self.vocab_size,)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name=F'out_layers_._{i}_._bias'   ,)
								self.out_layers.append((weight, bias)		)
						else:
							for i in range(len(self.cutoffs		)		):
								lowerCAmelCase__				,lowerCAmelCase__	:						Union[str, Any]  				=	self.cutoff_ends[i], self.cutoff_ends[i + 1]
								lowerCAmelCase__	:						int  				=	self.d_embed // (self.div_val**i)

								lowerCAmelCase__	:						Optional[Any]  				=	self.add_weight(
								    shape=(d_emb_i, self.d_proj)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name=F'out_projs_._{i}'		)
								self.out_projs.append(lowercase_		)
								lowerCAmelCase__	:						str  				=	self.add_weight(
								    shape=(r_idx - l_idx, d_emb_i)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name=F'out_layers_._{i}_._weight'   ,)
								lowerCAmelCase__	:						Optional[int]  				=	self.add_weight(
								    shape=(r_idx - l_idx,)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name=F'out_layers_._{i}_._bias'   ,)
								self.out_layers.append((weight, bias)		)
						super().build(lowercase_		)







					@staticmethod
					def        __lowerCAmelCase (       lowercase_			:				List[str]   ,lowercase_			:				int   ,lowercase_			:				Tuple   ,lowercase_			:				List[Any]=None		):
						lowerCAmelCase__	:						Optional[Any]  				=	x
						if proj is not None:
							lowerCAmelCase__	:						Optional[Any]  				=	tf.einsum('''ibd,ed->ibe'''   ,lowercase_   ,lowercase_		)
						return tf.einsum('''ibd,nd->ibn'''   ,lowercase_   ,lowercase_		) + b







					@staticmethod
					def        __lowerCAmelCase (       lowercase_			:				int   ,lowercase_			:				Tuple		):
						lowerCAmelCase__	:						Optional[int]  				=	shape_list(lowercase_		)
						lowerCAmelCase__	:						Dict  				=	tf.range(lp_size[0]   ,dtype=target.dtype		)
						lowerCAmelCase__	:						str  				=	tf.stack([r, target]   ,1		)
						return tf.gather_nd(lowercase_   ,lowercase_		)







					def        __lowerCAmelCase (       self			:				Optional[int]   ,lowercase_			:				str   ,lowercase_			:				Dict   ,lowercase_			:				Dict=True   ,lowercase_			:				Dict=False		):
						lowerCAmelCase__	:						Optional[Any]  				=	0
						if self.n_clusters == 0:
							lowerCAmelCase__	:						Union[str, Any]  				=	self._logit(lowercase_   ,self.out_layers[0][0]   ,self.out_layers[0][1]   ,self.out_projs[0]		)
							if target is not None:
								lowerCAmelCase__	:						Any  				=	tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowercase_   ,logits=lowercase_		)
							lowerCAmelCase__	:						Any  				=	tf.nn.log_softmax(lowercase_   ,axis=-1		)
						else:
							lowerCAmelCase__	:						Optional[Any]  				=	shape_list(lowercase_		)
							lowerCAmelCase__	:						Union[str, Any]  				=	[]
							lowerCAmelCase__	:						int  				=	tf.zeros(hidden_sizes[:2]		)
							for i in range(len(self.cutoffs		)		):
								lowerCAmelCase__				,lowerCAmelCase__	:						Any  				=	self.cutoff_ends[i], self.cutoff_ends[i + 1]
								if target is not None:
									lowerCAmelCase__	:						int  				=	(target >= l_idx) & (target < r_idx)
									lowerCAmelCase__	:						List[Any]  				=	tf.where(lowercase_		)
									lowerCAmelCase__	:						Union[str, Any]  				=	tf.boolean_mask(lowercase_   ,lowercase_		) - l_idx

								if self.div_val == 1:
									lowerCAmelCase__	:						int  				=	self.out_layers[0][0][l_idx:r_idx]
									lowerCAmelCase__	:						Union[str, Any]  				=	self.out_layers[0][1][l_idx:r_idx]
								else:
									lowerCAmelCase__	:						Optional[Any]  				=	self.out_layers[i][0]
									lowerCAmelCase__	:						Optional[int]  				=	self.out_layers[i][1]

								if i == 0:
									lowerCAmelCase__	:						Union[str, Any]  				=	tf.concat([cur_W, self.cluster_weight]   ,0		)
									lowerCAmelCase__	:						Tuple  				=	tf.concat([cur_b, self.cluster_bias]   ,0		)

									lowerCAmelCase__	:						List[Any]  				=	self._logit(lowercase_   ,lowercase_   ,lowercase_   ,self.out_projs[0]		)
									lowerCAmelCase__	:						List[str]  				=	tf.nn.log_softmax(lowercase_		)
									out.append(head_logprob[..., : self.cutoffs[0]]		)
									if target is not None:
										lowerCAmelCase__	:						List[str]  				=	tf.boolean_mask(lowercase_   ,lowercase_		)
										lowerCAmelCase__	:						Union[str, Any]  				=	self._gather_logprob(lowercase_   ,lowercase_		)
								else:
									lowerCAmelCase__	:						int  				=	self._logit(lowercase_   ,lowercase_   ,lowercase_   ,self.out_projs[i]		)
									lowerCAmelCase__	:						Tuple  				=	tf.nn.log_softmax(lowercase_		)
									lowerCAmelCase__	:						Tuple  				=	self.cutoffs[0] + i - 1  # No probability for the head cluster
									lowerCAmelCase__	:						Optional[int]  				=	head_logprob[..., cluster_prob_idx, None] + tail_logprob
									out.append(lowercase_		)
									if target is not None:
										lowerCAmelCase__	:						Optional[Any]  				=	tf.boolean_mask(lowercase_   ,lowercase_		)
										lowerCAmelCase__	:						Dict  				=	tf.boolean_mask(lowercase_   ,lowercase_		)
										lowerCAmelCase__	:						str  				=	self._gather_logprob(lowercase_   ,lowercase_		)
										cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
								if target is not None:
									loss += tf.scatter_nd(lowercase_   ,-cur_logprob   ,shape_list(lowercase_		)		)
							lowerCAmelCase__	:						Optional[int]  				=	tf.concat(lowercase_   ,axis=-1		)

						if target is not None:
							if return_mean:
								lowerCAmelCase__	:						List[Any]  				=	tf.reduce_mean(lowercase_		)
							# Add the training-time loss value to the layer using `self.add_loss()`.
							self.add_loss(lowercase_		)

							# Log the loss as a metric (we could log arbitrary metrics,
							# including different metrics for training and inference.
							self.add_metric(lowercase_   ,name=self.name   ,aggregation='''mean''' if return_mean else ''''''		)

						return out





"""simple docstring"""






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



    """simple docstring"""

    _UpperCAmelCase	:Union[str, Any]						    = MBartConfig
    _UpperCAmelCase	:str						    = {}
    _UpperCAmelCase	:Union[str, Any]						    = "gelu"






    def __init__(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase=13    ,			_UpperCAmelCase=7    ,			_UpperCAmelCase=True    ,			_UpperCAmelCase=False    ,			_UpperCAmelCase=99    ,			_UpperCAmelCase=32    ,			_UpperCAmelCase=2    ,			_UpperCAmelCase=4    ,			_UpperCAmelCase=37    ,			_UpperCAmelCase=0.1    ,			_UpperCAmelCase=0.1    ,			_UpperCAmelCase=20    ,			_UpperCAmelCase=2    ,			_UpperCAmelCase=1    ,			_UpperCAmelCase=0    ,			):
       lowercase__:   List[str]					=      parent
       lowercase__:   List[Any]					=      batch_size
       lowercase__:   List[Any]					=      seq_length
       lowercase__:   str					=      is_training
       lowercase__:   List[str]					=      use_labels
       lowercase__:   Optional[int]					=      vocab_size
       lowercase__:   int					=      hidden_size
       lowercase__:   Optional[int]					=      num_hidden_layers
       lowercase__:   int					=      num_attention_heads
       lowercase__:   Tuple					=      intermediate_size
       lowercase__:   Optional[Any]					=      hidden_dropout_prob
       lowercase__:   Optional[int]					=      attention_probs_dropout_prob
       lowercase__:   str					=      max_position_embeddings
       lowercase__:   Union[str, Any]					=      eos_token_id
       lowercase__:   int					=      pad_token_id
       lowercase__:   List[str]					=      bos_token_id






    def _snake_case			(   self							):
       lowercase__:   int					=      ids_tensor([self.batch_size, self.seq_length - 1]    ,			self.vocab_size							)
       lowercase__:   str					=      tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size							)    ,			1							)
       lowercase__:   List[Any]					=      tf.concat([input_ids, eos_tensor]    ,			axis=1							)

       lowercase__:   List[Any]					=      ids_tensor([self.batch_size, self.seq_length]    ,			self.vocab_size							)

       lowercase__:   Optional[Any]					=      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    ,			)
       lowercase__:   Optional[int]					=      prepare_mbart_inputs_dict(_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							)
       return config, inputs_dict






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
       lowercase__:   List[str]					=      TFMBartModel(config=_UpperCAmelCase							).get_decoder()
       lowercase__:   Tuple					=      inputs_dict['''input_ids''']

       lowercase__:   Optional[Any]					=      input_ids[:1, :]
       lowercase__:   Optional[int]					=      inputs_dict['''attention_mask'''][:1, :]
       lowercase__:   List[str]					=      inputs_dict['''head_mask''']
       lowercase__:   Optional[int]					=      1

       # first forward pass
       lowercase__:   List[Any]					=      model(_UpperCAmelCase    ,			attention_mask=_UpperCAmelCase    ,			head_mask=_UpperCAmelCase    ,			use_cache=_UpperCAmelCase							)

       lowercase__, lowercase__:   Any					=      outputs.to_tuple()
       lowercase__:   List[str]					=      past_key_values[1]



def SCREAMING_SNAKE_CASE__			(   __UpperCAmelCase       ,    __UpperCAmelCase       ,    __UpperCAmelCase       ,    __UpperCAmelCase=None       ,    __UpperCAmelCase=None       ,    __UpperCAmelCase=None       ,    __UpperCAmelCase=None       ,    __UpperCAmelCase=None       ,    )  ->  int:
   if attention_mask is None:
      lowercase__:   Union[str, Any]					=      tf.cast(tf.math.not_equal(__UpperCAmelCase       ,    config.pad_token_id )       ,    tf.inta )
   if decoder_attention_mask is None:
      lowercase__:   List[str]					=      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:
      lowercase__:   str					=      tf.ones((config.encoder_layers, config.encoder_attention_heads) )
   if decoder_head_mask is None:
      lowercase__:   List[str]					=      tf.ones((config.decoder_layers, config.decoder_attention_heads) )
   if cross_attn_head_mask is None:
      lowercase__:   Union[str, Any]					=      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     UpperCAmelCase						(_UpperCAmelCase    ,_UpperCAmelCase    ,unittest.TestCase  ):



    """simple docstring"""

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






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
       if pipeline_test_casse_name != "FeatureExtractionPipelineTests":
          # Exception encountered when calling layer '...'
          return True

       return False






    def _snake_case			(   self							):
       lowercase__:   Tuple					=      TFMBartModelTester(self							)
       lowercase__:   Optional[Any]					=      ConfigTester(self    ,			config_class=_UpperCAmelCase							)






    def _snake_case			(   self							):
       self.config_tester.run_common_tests()






    def _snake_case			(   self							):
       lowercase__:   List[str]					=      self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.check_decoder_model_past_large_inputs(*_UpperCAmelCase							)



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



    """simple docstring"""

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






    @cached_property
    def _snake_case			(   self							):
       return AutoTokenizer.from_pretrained(self.model_name							)






    @cached_property
    def _snake_case			(   self							):
       lowercase__:   Optional[int]					=      TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name							)
       return model






    def _snake_case			(   self    ,			**_UpperCAmelCase							):
       lowercase__:   List[Any]					=      self.translate_src_text(**_UpperCAmelCase							)
       self.assertListEqual(self.expected_text    ,			_UpperCAmelCase							)






    def _snake_case			(   self    ,			**_UpperCAmelCase							):
       lowercase__:   str					=      self.tokenizer(self.src_text    ,			**_UpperCAmelCase    ,			return_tensors='''tf'''							)
       lowercase__:   Optional[int]					=      self.model.generate(
           model_inputs.input_ids    ,			attention_mask=model_inputs.attention_mask    ,			num_beams=2							)
       lowercase__:   Tuple					=      self.tokenizer.batch_decode(_UpperCAmelCase    ,			skip_special_tokens=_UpperCAmelCase							)
       return generated_words






    @slow
    def _snake_case			(   self							):
       self._assert_generated_batch_equal_expected()
