




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  __magic_name__						:
				'''simple docstring'''







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


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


				def     _A							(							self:   Optional[int]		):
											SCREAMING_SNAKE_CASE_   	=    ids_tensor([self.batch_size, self.seq_length - 1]		,					self.vocab_size		)
											SCREAMING_SNAKE_CASE_   	=    tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size		)		,					1		)
											SCREAMING_SNAKE_CASE_   	=    tf.concat([input_ids, eos_tensor]		,					axis=1		)

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

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





				def     _A							(							self:   Tuple		,					_lowerCamelCase:   int		,					_lowerCamelCase:   List[Any]		):
											SCREAMING_SNAKE_CASE_   	=    TFMBartModel(config=_lowerCamelCase		).get_decoder()
											SCREAMING_SNAKE_CASE_   	=    inputs_dict['''input_ids''']

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

											# first forward pass
											SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		,					attention_mask=_lowerCamelCase		,					head_mask=_lowerCamelCase		,					use_cache=_lowerCamelCase		)

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






def 			a   (_lowerCAmelCase  ,    _lowerCAmelCase  ,    _lowerCAmelCase  ,    _lowerCAmelCase=None  ,    _lowerCAmelCase=None  ,    _lowerCAmelCase=None  ,    _lowerCAmelCase=None  ,    _lowerCAmelCase=None  ,    ):
							if attention_mask is None:
														SCREAMING_SNAKE_CASE_   	=    tf.cast(tf.math.not_equal(_lowerCAmelCase  ,    config.pad_token_id     )  ,    tf.inta     )
							if decoder_attention_mask is None:
														SCREAMING_SNAKE_CASE_   	=    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:
														SCREAMING_SNAKE_CASE_   	=    tf.ones((config.encoder_layers, config.encoder_attention_heads)     )
							if decoder_head_mask is None:
														SCREAMING_SNAKE_CASE_   	=    tf.ones((config.decoder_layers, config.decoder_attention_heads)     )
							if cross_attn_head_mask is None:
														SCREAMING_SNAKE_CASE_   	=    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  __magic_name__						(  __UpperCAmelCase    ,  __UpperCAmelCase    ,  unittest.TestCase):
				'''simple docstring'''







				SCREAMING_SNAKE_CASE__					:							Tuple	      =					(TFMBartForConditionalGeneration, TFMBartModel) if is_tf_available() else ()
				SCREAMING_SNAKE_CASE__					:							int	      =					(TFMBartForConditionalGeneration,) if is_tf_available() else ()
				SCREAMING_SNAKE_CASE__					:							Tuple	      =					(
				    {
				        "conversational": TFMBartForConditionalGeneration,
				        "feature-extraction": TFMBartModel,
				        "summarization": TFMBartForConditionalGeneration,
				        "text2text-generation": TFMBartForConditionalGeneration,
				        "translation": TFMBartForConditionalGeneration,
				    }
				    if is_tf_available()
				    else {}
				)
				SCREAMING_SNAKE_CASE__					:							Any	      =					True
				SCREAMING_SNAKE_CASE__					:							Optional[int]	      =					False
				SCREAMING_SNAKE_CASE__					:							Tuple	      =					False


				def     _A							(							self:   int		,					_lowerCamelCase:   Union[str, Any]		,					_lowerCamelCase:   Union[str, Any]		,					_lowerCamelCase:   List[str]		,					_lowerCamelCase:   Any		,					_lowerCamelCase:   str		):
											if pipeline_test_casse_name != "FeatureExtractionPipelineTests":
																		# Exception encountered when calling layer '...'
																		return True

											return False


				def     _A							(							self:   Optional[Any]		):
											SCREAMING_SNAKE_CASE_   	=    TFMBartModelTester(self		)
											SCREAMING_SNAKE_CASE_   	=    ConfigTester(self		,					config_class=_lowerCamelCase		)


				def     _A							(							self:   Tuple		):
											self.config_tester.run_common_tests()





				def     _A							(							self:   Optional[Any]		):
											SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs_for_common()
											self.model_tester.check_decoder_model_past_large_inputs(*_lowerCamelCase		)






@require_sentencepiece
@require_tokenizers
@require_tf
class  __magic_name__						(  unittest.TestCase):
				'''simple docstring'''







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


				@cached_property
				def     _A							(							self:   List[Any]		):
											return AutoTokenizer.from_pretrained(self.model_name		)


				@cached_property
				def     _A							(							self:   str		):
											SCREAMING_SNAKE_CASE_   	=    TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name		)
											return model


				def     _A							(							self:   Dict		,					**_lowerCamelCase:   Optional[Any]		):
											SCREAMING_SNAKE_CASE_   	=    self.translate_src_text(**_lowerCamelCase		)
											self.assertListEqual(self.expected_text		,					_lowerCamelCase		)


				def     _A							(							self:   Union[str, Any]		,					**_lowerCamelCase:   Optional[Any]		):
											SCREAMING_SNAKE_CASE_   	=    self.tokenizer(self.src_text		,					**_lowerCamelCase		,					return_tensors='''tf'''		)
											SCREAMING_SNAKE_CASE_   	=    self.model.generate(
											    model_inputs.input_ids		,					attention_mask=model_inputs.attention_mask		,					num_beams=2		)
											SCREAMING_SNAKE_CASE_   	=    self.tokenizer.batch_decode(_lowerCamelCase		,					skip_special_tokens=_lowerCamelCase		)
											return generated_words





				@slow
				def     _A							(							self:   List[str]		):
											self._assert_generated_batch_equal_expected()








import math
import time
from typing import Dict, List, Optional

from torch.utils.data import Dataset

from transformers import SeqaSeqTrainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
			import torch_xla.core.xla_model as xm
			import torch_xla.debug.metrics as met






class  __magic_name__						(  __UpperCAmelCase):
				'''simple docstring'''







				def __init__(							self:   List[Any]		,					*_lowerCamelCase:   Tuple		,					_lowerCamelCase:   Optional[Any]=None		,					_lowerCamelCase:   Dict=None		,					**_lowerCamelCase:   Dict		):
											super().__init__(*_lowerCamelCase		,					**_lowerCamelCase		)
											SCREAMING_SNAKE_CASE_   	=    eval_examples
											SCREAMING_SNAKE_CASE_   	=    post_process_function


				def     _A							(							self:   Tuple		,					_lowerCamelCase:   Optional[Dataset] = None		,					_lowerCamelCase:   List[str]=None		,					_lowerCamelCase:   Optional[List[str]] = None		,					_lowerCamelCase:   str = "eval"		,					**_lowerCamelCase:   Union[str, Any]		,					):
											SCREAMING_SNAKE_CASE_   	=    gen_kwargs.copy()
											SCREAMING_SNAKE_CASE_   	=    (
											    gen_kwargs['''max_length'''] if gen_kwargs.get('''max_length'''		) is not None else self.args.generation_max_length
											)
											SCREAMING_SNAKE_CASE_   	=    (
											    gen_kwargs['''num_beams'''] if gen_kwargs.get('''num_beams'''		) is not None else self.args.generation_num_beams
											)
											SCREAMING_SNAKE_CASE_   	=    gen_kwargs

											SCREAMING_SNAKE_CASE_   	=    self.eval_dataset if eval_dataset is None else eval_dataset
											SCREAMING_SNAKE_CASE_   	=    self.get_eval_dataloader(_lowerCamelCase		)
											SCREAMING_SNAKE_CASE_   	=    self.eval_examples if eval_examples is None else eval_examples

											# Temporarily disable metric computation, we will do it in the loop here.
											SCREAMING_SNAKE_CASE_   	=    self.compute_metrics
											SCREAMING_SNAKE_CASE_   	=    None
											SCREAMING_SNAKE_CASE_   	=    time.time()
											SCREAMING_SNAKE_CASE_   	=    self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
											try:
																		SCREAMING_SNAKE_CASE_   	=    eval_loop(
																		    _lowerCamelCase		,					description='''Evaluation'''		,					prediction_loss_only=True if compute_metrics is None else None		,					ignore_keys=_lowerCamelCase		,					metric_key_prefix=_lowerCamelCase		,					)
											finally:
																		SCREAMING_SNAKE_CASE_   	=    compute_metrics
											SCREAMING_SNAKE_CASE_   	=    self.args.eval_batch_size * self.args.world_size
											if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
																		start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
											output.metrics.update(
											    speed_metrics(
											        _lowerCamelCase		,					_lowerCamelCase		,					num_samples=output.num_samples		,					num_steps=math.ceil(output.num_samples / total_batch_size		)		,					)		)

											if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
																		# Only the main node write the results by default
																		SCREAMING_SNAKE_CASE_   	=    self.post_process_function(_lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		)
																		SCREAMING_SNAKE_CASE_   	=    self.compute_metrics(_lowerCamelCase		)

																		# Prefix all keys with metric_key_prefix + '_'
																		for key in list(metrics.keys()		):
																									if not key.startswith(f"{metric_key_prefix}_"		):
																																SCREAMING_SNAKE_CASE_   	=    metrics.pop(_lowerCamelCase		)

																		metrics.update(output.metrics		)
											else:
																		SCREAMING_SNAKE_CASE_   	=    output.metrics

											if self.args.should_log:
																		# Only the main node log the results by default
																		self.log(_lowerCamelCase		)

											if self.args.tpu_metrics_debug or self.args.debug:
																		# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
																		xm.master_print(met.metrics_report()		)

											SCREAMING_SNAKE_CASE_   	=    self.callback_handler.on_evaluate(self.args		,					self.state		,					self.control		,					_lowerCamelCase		)
											return metrics





				def     _A							(							self:   List[str]		,					_lowerCamelCase:   int		,					_lowerCamelCase:   List[str]		,					_lowerCamelCase:   Optional[int]=None		,					_lowerCamelCase:   str = "test"		,					**_lowerCamelCase:   str		):
											SCREAMING_SNAKE_CASE_   	=    gen_kwargs.copy()

											SCREAMING_SNAKE_CASE_   	=    self.get_test_dataloader(_lowerCamelCase		)

											# Temporarily disable metric computation, we will do it in the loop here.
											SCREAMING_SNAKE_CASE_   	=    self.compute_metrics
											SCREAMING_SNAKE_CASE_   	=    None
											SCREAMING_SNAKE_CASE_   	=    time.time()
											SCREAMING_SNAKE_CASE_   	=    self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
											try:
																		SCREAMING_SNAKE_CASE_   	=    eval_loop(
																		    _lowerCamelCase		,					description='''Prediction'''		,					prediction_loss_only=True if compute_metrics is None else None		,					ignore_keys=_lowerCamelCase		,					metric_key_prefix=_lowerCamelCase		,					)
											finally:
																		SCREAMING_SNAKE_CASE_   	=    compute_metrics

											SCREAMING_SNAKE_CASE_   	=    self.args.eval_batch_size * self.args.world_size
											if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
																		start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
											output.metrics.update(
											    speed_metrics(
											        _lowerCamelCase		,					_lowerCamelCase		,					num_samples=output.num_samples		,					num_steps=math.ceil(output.num_samples / total_batch_size		)		,					)		)
											if self.post_process_function is None or self.compute_metrics is None:
																		return output

											SCREAMING_SNAKE_CASE_   	=    self.post_process_function(_lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		,					'''predict'''		)
											SCREAMING_SNAKE_CASE_   	=    self.compute_metrics(_lowerCamelCase		)

											# Prefix all keys with metric_key_prefix + '_'
											for key in list(metrics.keys()		):
																		if not key.startswith(f"{metric_key_prefix}_"		):
																									SCREAMING_SNAKE_CASE_   	=    metrics.pop(_lowerCamelCase		)
											metrics.update(output.metrics		)
											return PredictionOutput(predictions=predictions.predictions		,					label_ids=predictions.label_ids		,					metrics=_lowerCamelCase		)


