


import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import MBartConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
from transformers.utils import cached_property

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
					import os

					# The slow tests are often failing with OOM error on GPU
					# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
					# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
					UpperCamelCase :				Optional[int]						    =       '''platform'''

					import jax
					import jax.numpy as jnp

					from transformers import AutoTokenizer
					from transformers.models.mbart.modeling_flax_mbart import (
					    FlaxMBartForConditionalGeneration,
					    FlaxMBartForQuestionAnswering,
					    FlaxMBartForSequenceClassification,
					    FlaxMBartModel,
					    shift_tokens_right,
					)
def 							UpperCamelCase_  (				__snake_case			:						Optional[int]	,			__snake_case			:						Tuple	,			__snake_case			:						str=None	,			__snake_case			:						Optional[Any]=None	,			__snake_case			:						Optional[Any]=None	,			__snake_case			:						Dict=None	,			__snake_case			:						Dict=None	,			__snake_case			:						str=None	,			)->      Union[str, Any]:




						'''simple docstring'''
						if attention_mask is None:
												__magic_name__ 		=			np.where(input_ids != config.pad_token_id	,			1	,			0     )
						if decoder_attention_mask is None:
												__magic_name__ 		=			np.where(decoder_input_ids != config.pad_token_id	,			1	,			0     )
						if head_mask is None:
												__magic_name__ 		=			np.ones((config.encoder_layers, config.encoder_attention_heads)     )
						if decoder_head_mask is None:
												__magic_name__ 		=			np.ones((config.decoder_layers, config.decoder_attention_heads)     )
						if cross_attn_head_mask is None:
												__magic_name__ 		=			np.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,
						}

class     __snake_case							(      unittest.TestCase):
					def __init__(	self					:	List[str]  ,			_lowerCAmelCase					:	Optional[int]  ,			_lowerCAmelCase					:	Tuple=1_3  ,			_lowerCAmelCase					:	Optional[int]=7  ,			_lowerCAmelCase					:	Tuple=True  ,			_lowerCAmelCase					:	int=False  ,			_lowerCAmelCase					:	List[Any]=9_9  ,			_lowerCAmelCase					:	List[str]=1_6  ,			_lowerCAmelCase					:	List[str]=2  ,			_lowerCAmelCase					:	str=4  ,			_lowerCAmelCase					:	Dict=4  ,			_lowerCAmelCase					:	List[Any]="gelu"  ,			_lowerCAmelCase					:	Tuple=0.1  ,			_lowerCAmelCase					:	Union[str, Any]=0.1  ,			_lowerCAmelCase					:	Any=3_2  ,			_lowerCAmelCase					:	List[str]=2  ,			_lowerCAmelCase					:	Any=1  ,			_lowerCAmelCase					:	Tuple=0  ,			_lowerCAmelCase					:	Optional[Any]=2  ,			_lowerCAmelCase					:	int=0.02  ,			)       ->				Any:



											'''simple docstring'''





											__magic_name__ 		=			parent
											__magic_name__ 		=			batch_size
											__magic_name__ 		=			seq_length
											__magic_name__ 		=			is_training
											__magic_name__ 		=			use_labels
											__magic_name__ 		=			vocab_size
											__magic_name__ 		=			hidden_size
											__magic_name__ 		=			num_hidden_layers
											__magic_name__ 		=			num_attention_heads
											__magic_name__ 		=			intermediate_size
											__magic_name__ 		=			hidden_act
											__magic_name__ 		=			hidden_dropout_prob
											__magic_name__ 		=			attention_probs_dropout_prob
											__magic_name__ 		=			max_position_embeddings
											__magic_name__ 		=			eos_token_id
											__magic_name__ 		=			pad_token_id
											__magic_name__ 		=			bos_token_id
											__magic_name__ 		=			decoder_start_token_id
											__magic_name__ 		=			initializer_range
					def   a (	self					:	Any     )       ->				Union[str, Any]:



											'''simple docstring'''





											__magic_name__ 		=			np.clip(ids_tensor([self.batch_size, self.seq_length - 1]  ,			self.vocab_size     )  ,			3  ,			self.vocab_size     )
											__magic_name__ 		=			np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)  ,			dtype=np.intaa     ))  ,			-1     )

											__magic_name__ 		=			shift_tokens_right(_lowerCAmelCase  ,			1     )

											__magic_name__ 		=			MBartConfig(
											    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_id=self.eos_token_id  ,			bos_token_id=self.bos_token_id  ,			pad_token_id=self.pad_token_id  ,			decoder_start_token_id=self.decoder_start_token_id  ,			initializer_range=self.initializer_range  ,			use_cache=_lowerCAmelCase  ,			)
											__magic_name__ 		=			prepare_mbart_inputs_dict(_lowerCAmelCase  ,			_lowerCAmelCase  ,			_lowerCAmelCase     )
											return config, inputs_dict
					def   a (	self					:	List[str]     )       ->				int:



											'''simple docstring'''





											__magic_name__		,			__magic_name__ 		=			self.prepare_config_and_inputs()
											return config, inputs_dict
					def   a (	self					:	Optional[Any]  ,			_lowerCAmelCase					:	Optional[Any]  ,			_lowerCAmelCase					:	Union[str, Any]  ,			_lowerCAmelCase					:	Any     )       ->				Tuple:



											'''simple docstring'''





											__magic_name__ 		=			2_0
											__magic_name__ 		=			model_class_name(_lowerCAmelCase     )

											__magic_name__ 		=			model.encode(inputs_dict['input_ids']     )

											__magic_name__		,			__magic_name__ 		=			(
											    inputs_dict['decoder_input_ids'],
											    inputs_dict['decoder_attention_mask'],
											)

											__magic_name__ 		=			model.init_cache(decoder_input_ids.shape[0]  ,			_lowerCAmelCase  ,			_lowerCAmelCase     )
											__magic_name__ 		=			jnp.ones((decoder_input_ids.shape[0], max_decoder_length)  ,			dtype='i4'     )

											__magic_name__ 		=			jnp.broadcast_to(
											    jnp.arange(decoder_input_ids.shape[-1] - 1     )[None, :]  ,			(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)  ,			)
											__magic_name__ 		=			model.decode(
											    decoder_input_ids[:, :-1]  ,			_lowerCAmelCase  ,			decoder_attention_mask=_lowerCAmelCase  ,			past_key_values=_lowerCAmelCase  ,			decoder_position_ids=_lowerCAmelCase  ,			)

											__magic_name__ 		=			jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]  ,			dtype='i4'     )
											__magic_name__ 		=			model.decode(
											    decoder_input_ids[:, -1:]  ,			_lowerCAmelCase  ,			decoder_attention_mask=_lowerCAmelCase  ,			past_key_values=outputs_cache.past_key_values  ,			decoder_position_ids=_lowerCAmelCase  ,			)

											__magic_name__ 		=			model.decode(_lowerCAmelCase  ,			_lowerCAmelCase     )

											__magic_name__ 		=			np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])     )     )
											self.parent.assertTrue(diff < 1E-3  ,			msg=F'''Max diff is {diff}'''     )







					def   a (	self					:	str  ,			_lowerCAmelCase					:	Tuple  ,			_lowerCAmelCase					:	List[str]  ,			_lowerCAmelCase					:	Dict     )       ->				List[Any]:



											'''simple docstring'''





											__magic_name__ 		=			2_0
											__magic_name__ 		=			model_class_name(_lowerCAmelCase     )

											__magic_name__ 		=			model.encode(inputs_dict['input_ids']     )

											__magic_name__		,			__magic_name__ 		=			(
											    inputs_dict['decoder_input_ids'],
											    inputs_dict['decoder_attention_mask'],
											)

											__magic_name__ 		=			jnp.concatenate(
											    [
											        decoder_attention_mask,
											        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])     ),
											    ]  ,			axis=-1  ,			)

											__magic_name__ 		=			model.init_cache(decoder_input_ids.shape[0]  ,			_lowerCAmelCase  ,			_lowerCAmelCase     )
											__magic_name__ 		=			jnp.broadcast_to(
											    jnp.arange(decoder_input_ids.shape[-1] - 1     )[None, :]  ,			(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)  ,			)

											__magic_name__ 		=			model.decode(
											    decoder_input_ids[:, :-1]  ,			_lowerCAmelCase  ,			decoder_attention_mask=_lowerCAmelCase  ,			past_key_values=_lowerCAmelCase  ,			decoder_position_ids=_lowerCAmelCase  ,			)
											__magic_name__ 		=			jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]  ,			dtype='i4'     )
											__magic_name__ 		=			model.decode(
											    decoder_input_ids[:, -1:]  ,			_lowerCAmelCase  ,			past_key_values=outputs_cache.past_key_values  ,			decoder_attention_mask=_lowerCAmelCase  ,			decoder_position_ids=_lowerCAmelCase  ,			)

											__magic_name__ 		=			model.decode(_lowerCAmelCase  ,			_lowerCAmelCase  ,			decoder_attention_mask=_lowerCAmelCase     )

											__magic_name__ 		=			np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])     )     )
											self.parent.assertTrue(diff < 1E-3  ,			msg=F'''Max diff is {diff}'''     )







@require_flax
class     __snake_case							(      unittest.TestCase):
					SCREAMING_SNAKE_CASE_		:					Any 				= 99
					def   a (	self					:	int     )       ->				Optional[int]:



											'''simple docstring'''





											__magic_name__ 		=			np.array(
											    [
											        [7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2],
											        [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2],
											        [5, 9_7, 1_7, 3_9, 9_4, 4_0, 2],
											        [7_6, 8_3, 9_4, 2_5, 7_0, 7_8, 2],
											        [8_7, 5_9, 4_1, 3_5, 4_8, 6_6, 2],
											        [5_5, 1_3, 1_6, 5_8, 5, 2, 1],  # note padding
											        [6_4, 2_7, 3_1, 5_1, 1_2, 7_5, 2],
											        [5_2, 6_4, 8_6, 1_7, 8_3, 3_9, 2],
											        [4_8, 6_1, 9, 2_4, 7_1, 8_2, 2],
											        [2_6, 1, 6_0, 4_8, 2_2, 1_3, 2],
											        [2_1, 5, 6_2, 2_8, 1_4, 7_6, 2],
											        [4_5, 9_8, 3_7, 8_6, 5_9, 4_8, 2],
											        [7_0, 7_0, 5_0, 9, 2_8, 0, 2],
											    ]  ,			dtype=np.intaa  ,			)

											__magic_name__ 		=			input_ids.shape[0]
											__magic_name__ 		=			MBartConfig(
											    vocab_size=self.vocab_size  ,			d_model=2_4  ,			encoder_layers=2  ,			decoder_layers=2  ,			encoder_attention_heads=2  ,			decoder_attention_heads=2  ,			encoder_ffn_dim=3_2  ,			decoder_ffn_dim=3_2  ,			max_position_embeddings=4_8  ,			eos_token_id=2  ,			pad_token_id=1  ,			bos_token_id=0  ,			)
											return config, input_ids, batch_size
					def   a (	self					:	List[str]     )       ->				Union[str, Any]:



											'''simple docstring'''





											__magic_name__		,			__magic_name__		,			__magic_name__ 		=			self._get_config_and_data()
											__magic_name__ 		=			FlaxMBartForSequenceClassification(_lowerCAmelCase     )
											__magic_name__ 		=			model(input_ids=_lowerCAmelCase  ,			decoder_input_ids=_lowerCAmelCase     )
											__magic_name__ 		=			(batch_size, config.num_labels)
											self.assertEqual(outputs['logits'].shape  ,			_lowerCAmelCase     )
					def   a (	self					:	Dict     )       ->				List[str]:



											'''simple docstring'''





											__magic_name__		,			__magic_name__		,			__magic_name__ 		=			self._get_config_and_data()
											__magic_name__ 		=			FlaxMBartForQuestionAnswering(_lowerCAmelCase     )
											__magic_name__ 		=			model(input_ids=_lowerCAmelCase     )

											self.assertEqual(outputs['start_logits'].shape  ,			input_ids.shape     )
											self.assertEqual(outputs['end_logits'].shape  ,			input_ids.shape     )
					def   a (	self					:	Any     )       ->				Optional[int]:



											'''simple docstring'''





											__magic_name__		,			__magic_name__		,			__magic_name__ 		=			self._get_config_and_data()
											__magic_name__ 		=			FlaxMBartForConditionalGeneration(_lowerCAmelCase     )
											__magic_name__ 		=			lm_model(input_ids=_lowerCAmelCase     )
											__magic_name__ 		=			(batch_size, input_ids.shape[1], config.vocab_size)
											self.assertEqual(outputs['logits'].shape  ,			_lowerCAmelCase     )
					def   a (	self					:	str     )       ->				Optional[Any]:



											'''simple docstring'''





											__magic_name__ 		=			MBartConfig(
											    vocab_size=self.vocab_size  ,			d_model=1_4  ,			encoder_layers=2  ,			decoder_layers=2  ,			encoder_attention_heads=2  ,			decoder_attention_heads=2  ,			encoder_ffn_dim=8  ,			decoder_ffn_dim=8  ,			max_position_embeddings=4_8  ,			)
											__magic_name__ 		=			FlaxMBartForConditionalGeneration(_lowerCAmelCase     )
											__magic_name__ 		=			np.array([[7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2], [6_8, 3_4, 2_6, 5_8, 3_0, 2, 1]]  ,			dtype=np.intaa     )
											__magic_name__ 		=			np.array([[8_2, 7_1, 8_2, 1_8, 2], [5_8, 6_8, 2, 1, 1]]  ,			dtype=np.intaa     )
											__magic_name__ 		=			lm_model(input_ids=_lowerCAmelCase  ,			decoder_input_ids=_lowerCAmelCase     )
											__magic_name__ 		=			(*summary.shape, config.vocab_size)
											self.assertEqual(outputs['logits'].shape  ,			_lowerCAmelCase     )







					def   a (	self					:	Any     )       ->				Dict:



											'''simple docstring'''





											__magic_name__ 		=			np.array([[7_1, 8_2, 1_8, 3_3, 2, 1, 1], [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2]]  ,			dtype=np.intaa     )
											__magic_name__ 		=			shift_tokens_right(_lowerCAmelCase  ,			1     )
											__magic_name__ 		=			np.equal(_lowerCAmelCase  ,			1     ).astype(np.floataa     ).sum()
											__magic_name__ 		=			np.equal(_lowerCAmelCase  ,			1     ).astype(np.floataa     ).sum()
											self.assertEqual(shifted.shape  ,			input_ids.shape     )
											self.assertEqual(_lowerCAmelCase  ,			n_pad_before - 1     )
											self.assertTrue(np.equal(shifted[:, 0]  ,			2     ).all()     )







@require_flax
class     __snake_case							(      _UpperCAmelCase  ,			unittest.TestCase  ,			_UpperCAmelCase):
					SCREAMING_SNAKE_CASE_		:					Dict 				= True
					SCREAMING_SNAKE_CASE_		:					Tuple 				= (
					    (
					        FlaxMBartModel,
					        FlaxMBartForConditionalGeneration,
					        FlaxMBartForSequenceClassification,
					        FlaxMBartForQuestionAnswering,
					    )
					    if is_flax_available()
					    else ()
					)
					SCREAMING_SNAKE_CASE_		:					List[str] 				= (FlaxMBartForConditionalGeneration,) if is_flax_available() else ()
					def   a (	self					:	List[Any]     )       ->				Tuple:



											'''simple docstring'''





											__magic_name__ 		=			FlaxMBartModelTester(self     )
					def   a (	self					:	List[Any]     )       ->				Any:



											'''simple docstring'''





											__magic_name__		,			__magic_name__ 		=			self.model_tester.prepare_config_and_inputs()
											for model_class in self.all_model_classes:
																	self.model_tester.check_use_cache_forward(_lowerCAmelCase  ,			_lowerCAmelCase  ,			_lowerCAmelCase     )
					def   a (	self					:	Any     )       ->				Optional[int]:



											'''simple docstring'''





											__magic_name__		,			__magic_name__ 		=			self.model_tester.prepare_config_and_inputs()
											for model_class in self.all_model_classes:
																	self.model_tester.check_use_cache_forward_with_attn_mask(_lowerCAmelCase  ,			_lowerCAmelCase  ,			_lowerCAmelCase     )
					def   a (	self					:	List[Any]     )       ->				List[Any]:



											'''simple docstring'''





											__magic_name__		,			__magic_name__ 		=			self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	with self.subTest(model_class.__name__     ):
																							__magic_name__ 		=			self._prepare_for_class(_lowerCAmelCase  ,			_lowerCAmelCase     )
																							__magic_name__ 		=			model_class(_lowerCAmelCase     )

																							@jax.jit
																							def encode_jitted(_lowerCAmelCase					:	Tuple  ,			_lowerCAmelCase					:	Union[str, Any]=None  ,			**_lowerCAmelCase					:	List[Any]     ):
																													return model.encode(input_ids=_lowerCAmelCase  ,			attention_mask=_lowerCAmelCase     )

																							with self.subTest('JIT Enabled'     ):
																													__magic_name__ 		=			encode_jitted(**_lowerCAmelCase     ).to_tuple()

																							with self.subTest('JIT Disabled'     ):
																													with jax.disable_jit():
																																			__magic_name__ 		=			encode_jitted(**_lowerCAmelCase     ).to_tuple()

																							self.assertEqual(len(_lowerCAmelCase     )  ,			len(_lowerCAmelCase     )     )
																							for jitted_output, output in zip(_lowerCAmelCase  ,			_lowerCAmelCase     ):
																													self.assertEqual(jitted_output.shape  ,			output.shape     )
					def   a (	self					:	List[str]     )       ->				Tuple:



											'''simple docstring'''





											__magic_name__		,			__magic_name__ 		=			self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	with self.subTest(model_class.__name__     ):
																							__magic_name__ 		=			model_class(_lowerCAmelCase     )
																							__magic_name__ 		=			model.encode(inputs_dict['input_ids']  ,			inputs_dict['attention_mask']     )

																							__magic_name__ 		=			{
																							    'decoder_input_ids': inputs_dict['decoder_input_ids'],
																							    'decoder_attention_mask': inputs_dict['decoder_attention_mask'],
																							    'encoder_outputs': encoder_outputs,
																							}

																							@jax.jit
																							def decode_jitted(_lowerCAmelCase					:	Tuple  ,			_lowerCAmelCase					:	Tuple  ,			_lowerCAmelCase					:	int     ):
																													return model.decode(
																													    decoder_input_ids=_lowerCAmelCase  ,			decoder_attention_mask=_lowerCAmelCase  ,			encoder_outputs=_lowerCAmelCase  ,			)

																							with self.subTest('JIT Enabled'     ):
																													__magic_name__ 		=			decode_jitted(**_lowerCAmelCase     ).to_tuple()

																							with self.subTest('JIT Disabled'     ):
																													with jax.disable_jit():
																																			__magic_name__ 		=			decode_jitted(**_lowerCAmelCase     ).to_tuple()

																							self.assertEqual(len(_lowerCAmelCase     )  ,			len(_lowerCAmelCase     )     )
																							for jitted_output, output in zip(_lowerCAmelCase  ,			_lowerCAmelCase     ):
																													self.assertEqual(jitted_output.shape  ,			output.shape     )







					@slow
					def   a (	self					:	List[Any]     )       ->				int:



											'''simple docstring'''





											for model_class_name in self.all_model_classes:
																	__magic_name__ 		=			model_class_name.from_pretrained('facebook/mbart-large-cc25'  ,			from_pt=_lowerCAmelCase     )
																	# FlaxMBartForSequenceClassification expects eos token in input_ids
																	__magic_name__ 		=			np.ones((1, 1)     ) * model.config.eos_token_id
																	__magic_name__ 		=			model(_lowerCAmelCase     )
																	self.assertIsNotNone(_lowerCAmelCase     )







@require_flax
@require_sentencepiece
@require_tokenizers
class     __snake_case							(      unittest.TestCase):
					SCREAMING_SNAKE_CASE_		:					Optional[int] 				= [
					    " UN Chief Says There Is No Military Solution in Syria",
					]
					SCREAMING_SNAKE_CASE_		:					Any 				= [
					    "Şeful ONU declară că nu există o soluţie militară în Siria",
					]
					SCREAMING_SNAKE_CASE_		:					Optional[Any] 				= "facebook/mbart-large-en-ro"
					@cached_property
					def   a (	self					:	int     )       ->				Optional[int]:



											'''simple docstring'''





											return AutoTokenizer.from_pretrained(self.model_name     )
					@cached_property
					def   a (	self					:	Any     )       ->				Tuple:



											'''simple docstring'''





											__magic_name__ 		=			FlaxMBartForConditionalGeneration.from_pretrained(self.model_name  ,			from_pt=_lowerCAmelCase     )
											return model
					def   a (	self					:	Union[str, Any]  ,			**_lowerCAmelCase					:	Optional[int]     )       ->				Optional[int]:



											'''simple docstring'''





											__magic_name__ 		=			self.translate_src_text(**_lowerCAmelCase     )
											self.assertListEqual(self.expected_text  ,			_lowerCAmelCase     )
					def   a (	self					:	int  ,			**_lowerCAmelCase					:	Optional[int]     )       ->				Optional[int]:



											'''simple docstring'''





											__magic_name__ 		=			self.tokenizer(self.src_text  ,			**_lowerCAmelCase  ,			return_tensors='np'     )
											__magic_name__ 		=			self.model.generate(
											    model_inputs.input_ids  ,			attention_mask=model_inputs.attention_mask  ,			decoder_start_token_id=self.tokenizer.lang_code_to_id['ro_RO']  ,			early_stopping=_lowerCAmelCase  ,			num_beams=2  ,			).sequences
											__magic_name__ 		=			self.tokenizer.batch_decode(_lowerCAmelCase  ,			skip_special_tokens=_lowerCAmelCase     )
											return generated_words







					@slow
					def   a (	self					:	List[str]     )       ->				int:



											'''simple docstring'''





											self._assert_generated_batch_equal_expected()





def 							UpperCamelCase_  (				__snake_case			:						int = 1_00     )->      int:




						'''simple docstring'''
						__magic_name__ 		=			1
						__magic_name__ 		=			0
						for i in range(1	,			num + 1     ):
												fact *= i

						for j in str(__snake_case     ):
												result += int(__snake_case     )

						return result


if __name__ == "__main__":
					print(solution(int(input('''Enter the Number: ''').strip())))

