


"""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    __a						:
				lowerCamelCase     :       List[str]			       =MBartConfig
				lowerCamelCase     :       str			       ={}
				lowerCamelCase     :       int			       ='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      ,	):

					'''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      lowerCamelCase_     ( 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(UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase   )
					return config, inputs_dict







				def      lowerCamelCase_     ( self      ,	UpperCAmelCase      ,	UpperCAmelCase   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			TFMBartModel(config=UpperCAmelCase   ).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(UpperCAmelCase      ,	attention_mask=UpperCAmelCase      ,	head_mask=UpperCAmelCase      ,	use_cache=UpperCAmelCase   )

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




def      UpperCAmelCase     (       _lowercase :   Optional[int]			,   _lowercase :   List[str]			,   _lowercase :   Any			,   _lowercase :   Optional[Any]=None			,   _lowercase :   Any=None			,   _lowercase :   List[str]=None			,   _lowercase :   Optional[int]=None			,   _lowercase :   Any=None			,   )						-> Optional[Any]:


	"""simple docstring"""




	if attention_mask is None:
		lowerCAmelCase_      							=			tf.cast(tf.math.not_equal(_lowercase			,   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							):
				lowerCamelCase     :       Union[str, Any]			       =(TFMBartForConditionalGeneration, TFMBartModel) if is_tf_available() else ()
				lowerCamelCase     :       Optional[int]			       =(TFMBartForConditionalGeneration,) if is_tf_available() else ()
				lowerCamelCase     :       Optional[int]			       =(
				    {
				        'conversational': TFMBartForConditionalGeneration,
				        'feature-extraction': TFMBartModel,
				        'summarization': TFMBartForConditionalGeneration,
				        'text2text-generation': TFMBartForConditionalGeneration,
				        'translation': TFMBartForConditionalGeneration,
				    }
				    if is_tf_available()
				    else {}
				)
				lowerCamelCase     :       int			       =True
				lowerCamelCase     :       str			       =False
				lowerCamelCase     :       Dict			       =False







				def      lowerCamelCase_     ( self      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase   ):

					'''simple docstring'''






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

					return False







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			TFMBartModelTester(self   )
					lowerCAmelCase_      							=			ConfigTester(self      ,	config_class=UpperCAmelCase   )







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					self.config_tester.run_common_tests()







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			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    __a						( unittest.TestCase							):
				lowerCamelCase     :       int			       =[
				    ' 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     :       Dict			       ='facebook/mbart-large-en-ro'







				@cached_property
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					return AutoTokenizer.from_pretrained(self.model_name   )







				@cached_property
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






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







				def      lowerCamelCase_     ( self      ,	**UpperCAmelCase   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			self.translate_src_text(**UpperCAmelCase   )
					self.assertListEqual(self.expected_text      ,	UpperCAmelCase   )







				def      lowerCamelCase_     ( self      ,	**UpperCAmelCase   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			self.tokenizer(self.src_text      ,	**UpperCAmelCase      ,	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(UpperCAmelCase      ,	skip_special_tokens=UpperCAmelCase   )
					return generated_words







				@slow
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					self._assert_generated_batch_equal_expected()



"""simple docstring"""



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


lowercase_							=							logging.get_logger(__name__)

lowercase_							=							{
    's-JoL/Open-Llama-V1': 'https://huggingface.co/s-JoL/Open-Llama-V1/blob/main/config.json',
}







class    __a						( __snake_case							):
				lowerCamelCase     :       Union[str, Any]			       ='open-llama'







				def __init__( self      ,	UpperCAmelCase=10_0000      ,	UpperCAmelCase=4096      ,	UpperCAmelCase=1_1008      ,	UpperCAmelCase=32      ,	UpperCAmelCase=32      ,	UpperCAmelCase="silu"      ,	UpperCAmelCase=2048      ,	UpperCAmelCase=0.0_2      ,	UpperCAmelCase=1E-6      ,	UpperCAmelCase=True      ,	UpperCAmelCase=0      ,	UpperCAmelCase=1      ,	UpperCAmelCase=2      ,	UpperCAmelCase=False      ,	UpperCAmelCase=True      ,	UpperCAmelCase=0.1      ,	UpperCAmelCase=0.1      ,	UpperCAmelCase=True      ,	UpperCAmelCase=True      ,	UpperCAmelCase=None      ,	**UpperCAmelCase      ,	):

					'''simple docstring'''






					lowerCAmelCase_      							=			vocab_size
					lowerCAmelCase_      							=			max_position_embeddings
					lowerCAmelCase_      							=			hidden_size
					lowerCAmelCase_      							=			intermediate_size
					lowerCAmelCase_      							=			num_hidden_layers
					lowerCAmelCase_      							=			num_attention_heads
					lowerCAmelCase_      							=			hidden_act
					lowerCAmelCase_      							=			initializer_range
					lowerCAmelCase_      							=			rms_norm_eps
					lowerCAmelCase_      							=			use_cache
					lowerCAmelCase_      							=			kwargs.pop(
					    '''use_memorry_efficient_attention'''      ,	UpperCAmelCase   )
					lowerCAmelCase_      							=			hidden_dropout_prob
					lowerCAmelCase_      							=			attention_dropout_prob
					lowerCAmelCase_      							=			use_stable_embedding
					lowerCAmelCase_      							=			shared_input_output_embedding
					lowerCAmelCase_      							=			rope_scaling
					self._rope_scaling_validation()

					super().__init__(
					    pad_token_id=UpperCAmelCase      ,	bos_token_id=UpperCAmelCase      ,	eos_token_id=UpperCAmelCase      ,	tie_word_embeddings=UpperCAmelCase      ,	**UpperCAmelCase      ,	)







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					if self.rope_scaling is None:
						return

					if not isinstance(self.rope_scaling      ,	UpperCAmelCase   ) or len(self.rope_scaling   ) != 2:
						raise ValueError(
						    '''`rope_scaling` must be a dictionary with with two fields, `name` and `factor`, '''
						    F"""got {self.rope_scaling}"""   )
					lowerCAmelCase_      							=			self.rope_scaling.get('''type'''      ,	UpperCAmelCase   )
					lowerCAmelCase_      							=			self.rope_scaling.get('''factor'''      ,	UpperCAmelCase   )
					if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
						raise ValueError(
						    F"""`rope_scaling`'s name field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"""   )
					if rope_scaling_factor is None or not isinstance(UpperCAmelCase      ,	UpperCAmelCase   ) or rope_scaling_factor <= 1.0:
						raise ValueError(F"""`rope_scaling`'s factor field must be an float > 1, got {rope_scaling_factor}"""   )