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

							"""simple docstring"""






							_snake_case   						=     MBartConfig
							_snake_case   						=     {}
							_snake_case   						=     'gelu'



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





													'''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 			A__						(     self	)->      Optional[int]:





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



							def 			A__						(     self   ,		SCREAMING_SNAKE_CASE_   ,		SCREAMING_SNAKE_CASE_	)->      Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase        =	TFMBartModel(config=__SCREAMING_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(__SCREAMING_SNAKE_CASE   ,		attention_mask=__SCREAMING_SNAKE_CASE   ,		head_mask=__SCREAMING_SNAKE_CASE   ,		use_cache=__SCREAMING_SNAKE_CASE	)

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





def 		A_       (							snake_case							:      Optional[int]       ,					snake_case							:      Any       ,					snake_case							:      str       ,					snake_case							:      List[Any]=None       ,					snake_case							:      Dict=None       ,					snake_case							:      Union[str, Any]=None       ,					snake_case							:      int=None       ,					snake_case							:      List[str]=None       ,					)					->			List[str]:



						'''simple docstring'''

						if attention_mask is None:
												__UpperCamelCase        =	tf.cast(tf.math.not_equal(__SCREAMING_SNAKE_CASE       ,					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 				SCREAMING_SNAKE_CASE__			(							lowerCamelCase_      ,						lowerCamelCase_      ,						unittest.TestCase			):

							"""simple docstring"""






							_snake_case   						=     (TFMBartForConditionalGeneration, TFMBartModel) if is_tf_available() else ()
							_snake_case   						=     (TFMBartForConditionalGeneration,) if is_tf_available() else ()
							_snake_case   						=     (
							    {
							        'conversational': TFMBartForConditionalGeneration,
							        'feature-extraction': TFMBartModel,
							        'summarization': TFMBartForConditionalGeneration,
							        'text2text-generation': TFMBartForConditionalGeneration,
							        'translation': TFMBartForConditionalGeneration,
							    }
							    if is_tf_available()
							    else {}
							)
							_snake_case   						=     True
							_snake_case   						=     False
							_snake_case   						=     False



							def 			A__						(     self   ,		SCREAMING_SNAKE_CASE_   ,		SCREAMING_SNAKE_CASE_   ,		SCREAMING_SNAKE_CASE_   ,		SCREAMING_SNAKE_CASE_   ,		SCREAMING_SNAKE_CASE_	)->      int:





													'''simple docstring'''






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

													return False



							def 			A__						(     self	)->      Tuple:





													'''simple docstring'''






													__UpperCamelCase        =	TFMBartModelTester(self	)
													__UpperCamelCase        =	ConfigTester(self   ,		config_class=__SCREAMING_SNAKE_CASE	)



							def 			A__						(     self	)->      List[str]:





													'''simple docstring'''






													self.config_tester.run_common_tests()



							def 			A__						(     self	)->      Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase        =	self.model_tester.prepare_config_and_inputs_for_common()
													self.model_tester.check_decoder_model_past_large_inputs(*__SCREAMING_SNAKE_CASE	)






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

							"""simple docstring"""






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



							@cached_property
							def 			A__						(     self	)->      Tuple:





													'''simple docstring'''






													return AutoTokenizer.from_pretrained(self.model_name	)



							@cached_property
							def 			A__						(     self	)->      List[str]:





													'''simple docstring'''






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



							def 			A__						(     self   ,		**SCREAMING_SNAKE_CASE_	)->      Any:





													'''simple docstring'''






													__UpperCamelCase        =	self.translate_src_text(**__SCREAMING_SNAKE_CASE	)
													self.assertListEqual(self.expected_text   ,		__SCREAMING_SNAKE_CASE	)



							def 			A__						(     self   ,		**SCREAMING_SNAKE_CASE_	)->      Optional[Any]:





													'''simple docstring'''






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



							@slow
							def 			A__						(     self	)->      Tuple:





													'''simple docstring'''






													self._assert_generated_batch_equal_expected()






'''simple docstring'''







import json
import os
import tempfile

from transformers.testing_utils import check_json_file_has_correct_format






class  lowerCAmelCase__   :
					lowerCAmelCase_        =     None






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Tuple													=       self.feature_extraction_class(**self.feat_extract_dict		)
											lowercase_	:     Any													=       json.loads(feat_extract.to_json_string()		)
											for key, value in self.feat_extract_dict.items():
																	self.assertEqual(obj[key]				,							__SCREAMING_SNAKE_CASE		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     str													=       self.feature_extraction_class(**self.feat_extract_dict		)

											with tempfile.TemporaryDirectory() as tmpdirname:
																	lowercase_	:     str													=       os.path.join(__SCREAMING_SNAKE_CASE				,							'''feat_extract.json'''		)
																	feat_extract_first.to_json_file(__SCREAMING_SNAKE_CASE		)
																	lowercase_	:     str													=       self.feature_extraction_class.from_json_file(__SCREAMING_SNAKE_CASE		)

											self.assertEqual(feat_extract_second.to_dict()				,							feat_extract_first.to_dict()		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Tuple													=       self.feature_extraction_class(**self.feat_extract_dict		)

											with tempfile.TemporaryDirectory() as tmpdirname:
																	lowercase_	:     Union[str, Any]													=       feat_extract_first.save_pretrained(__SCREAMING_SNAKE_CASE		)[0]
																	check_json_file_has_correct_format(__SCREAMING_SNAKE_CASE		)
																	lowercase_	:     str													=       self.feature_extraction_class.from_pretrained(__SCREAMING_SNAKE_CASE		)

											self.assertEqual(feat_extract_second.to_dict()				,							feat_extract_first.to_dict()		)






					def 	_snake_case    (   self		):


											"""simple docstring"""





											lowercase_	:     Optional[Any]													=       self.feature_extraction_class()
											self.assertIsNotNone(__SCREAMING_SNAKE_CASE		)



