
"""simple docstring"""


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
			__A       						=       '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     _lowerCamelCase(__UpperCamelCase    , __UpperCamelCase    , __UpperCamelCase=None    , __UpperCamelCase=None    , __UpperCamelCase=None    , __UpperCamelCase=None    , __UpperCamelCase=None    , __UpperCamelCase=None    , )				-> Tuple:
	if attention_mask is None:
		_lowerCAmelCase								=np.where(input_ids != config.pad_token_id    , 1    , 0       )
	if decoder_attention_mask is None:
		_lowerCAmelCase								=np.where(decoder_input_ids != config.pad_token_id    , 1    , 0       )
	if head_mask is None:
		_lowerCAmelCase								=np.ones((config.encoder_layers, config.encoder_attention_heads)       )
	if decoder_head_mask is None:
		_lowerCAmelCase								=np.ones((config.decoder_layers, config.decoder_attention_heads)       )
	if cross_attn_head_mask is None:
		_lowerCAmelCase								=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 		lowerCamelCase__					(					unittest.TestCase				):





			'''simple docstring'''



			def __init__(							self     ,     __UpperCAmelCase     ,     __UpperCAmelCase=13     ,     __UpperCAmelCase=7     ,     __UpperCAmelCase=True     ,     __UpperCAmelCase=False     ,     __UpperCAmelCase=99     ,     __UpperCAmelCase=16     ,     __UpperCAmelCase=2     ,     __UpperCAmelCase=4     ,     __UpperCAmelCase=4     ,     __UpperCAmelCase="gelu"     ,     __UpperCAmelCase=0.1     ,     __UpperCAmelCase=0.1     ,     __UpperCAmelCase=32     ,     __UpperCAmelCase=2     ,     __UpperCAmelCase=1     ,     __UpperCAmelCase=0     ,     __UpperCAmelCase=2     ,     __UpperCAmelCase=0.0_2     ,     ) ->     Optional[Any]:
				_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_act
				_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
				_lowerCAmelCase								=decoder_start_token_id
				_lowerCAmelCase								=initializer_range



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

				_lowerCAmelCase								=shift_tokens_right(__UpperCAmelCase     ,     1  )

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



			def 					_lowerCAmelCase		(							self  ) ->     Any:
				_lowerCAmelCase						,							_lowerCAmelCase								=self.prepare_config_and_inputs()
				return config, inputs_dict



			def 					_lowerCAmelCase		(							self     ,     __UpperCAmelCase     ,     __UpperCAmelCase     ,     __UpperCAmelCase  ) ->     Optional[int]:
				_lowerCAmelCase								=20
				_lowerCAmelCase								=model_class_name(__UpperCAmelCase  )

				_lowerCAmelCase								=model.encode(inputs_dict["""input_ids"""]  )

				_lowerCAmelCase						,							_lowerCAmelCase								=(
				    inputs_dict["""decoder_input_ids"""],
				    inputs_dict["""decoder_attention_mask"""],
				)

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

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

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

				_lowerCAmelCase								=model.decode(__UpperCAmelCase     ,     __UpperCAmelCase  )

				_lowerCAmelCase								=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 					_lowerCAmelCase		(							self     ,     __UpperCAmelCase     ,     __UpperCAmelCase     ,     __UpperCAmelCase  ) ->     Optional[int]:
				_lowerCAmelCase								=20
				_lowerCAmelCase								=model_class_name(__UpperCAmelCase  )

				_lowerCAmelCase								=model.encode(inputs_dict["""input_ids"""]  )

				_lowerCAmelCase						,							_lowerCAmelCase								=(
				    inputs_dict["""decoder_input_ids"""],
				    inputs_dict["""decoder_attention_mask"""],
				)

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

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

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

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

				_lowerCAmelCase								=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 		lowerCamelCase__					(					unittest.TestCase				):





			'''simple docstring'''
			lowerCamelCase  					=		99



			def 					_lowerCAmelCase		(							self  ) ->     List[str]:
				_lowerCAmelCase								=np.array(
				    [
				        [71, 82, 18, 33, 46, 91, 2],
				        [68, 34, 26, 58, 30, 82, 2],
				        [5, 97, 17, 39, 94, 40, 2],
				        [76, 83, 94, 25, 70, 78, 2],
				        [87, 59, 41, 35, 48, 66, 2],
				        [55, 13, 16, 58, 5, 2, 1],  # note padding
				        [64, 27, 31, 51, 12, 75, 2],
				        [52, 64, 86, 17, 83, 39, 2],
				        [48, 61, 9, 24, 71, 82, 2],
				        [26, 1, 60, 48, 22, 13, 2],
				        [21, 5, 62, 28, 14, 76, 2],
				        [45, 98, 37, 86, 59, 48, 2],
				        [70, 70, 50, 9, 28, 0, 2],
				    ]     ,     dtype=np.intaa     ,     )

				_lowerCAmelCase								=input_ids.shape[0]
				_lowerCAmelCase								=MBartConfig(
				    vocab_size=self.vocab_size     ,     d_model=24     ,     encoder_layers=2     ,     decoder_layers=2     ,     encoder_attention_heads=2     ,     decoder_attention_heads=2     ,     encoder_ffn_dim=32     ,     decoder_ffn_dim=32     ,     max_position_embeddings=48     ,     eos_token_id=2     ,     pad_token_id=1     ,     bos_token_id=0     ,     )
				return config, input_ids, batch_size



			def 					_lowerCAmelCase		(							self  ) ->     str:
				_lowerCAmelCase						,							_lowerCAmelCase						,							_lowerCAmelCase								=self._get_config_and_data()
				_lowerCAmelCase								=FlaxMBartForSequenceClassification(__UpperCAmelCase  )
				_lowerCAmelCase								=model(input_ids=__UpperCAmelCase     ,     decoder_input_ids=__UpperCAmelCase  )
				_lowerCAmelCase								=(batch_size, config.num_labels)
				self.assertEqual(outputs["""logits"""].shape     ,     __UpperCAmelCase  )



			def 					_lowerCAmelCase		(							self  ) ->     Any:
				_lowerCAmelCase						,							_lowerCAmelCase						,							_lowerCAmelCase								=self._get_config_and_data()
				_lowerCAmelCase								=FlaxMBartForQuestionAnswering(__UpperCAmelCase  )
				_lowerCAmelCase								=model(input_ids=__UpperCAmelCase  )

				self.assertEqual(outputs["""start_logits"""].shape     ,     input_ids.shape  )
				self.assertEqual(outputs["""end_logits"""].shape     ,     input_ids.shape  )



			def 					_lowerCAmelCase		(							self  ) ->     Dict:
				_lowerCAmelCase						,							_lowerCAmelCase						,							_lowerCAmelCase								=self._get_config_and_data()
				_lowerCAmelCase								=FlaxMBartForConditionalGeneration(__UpperCAmelCase  )
				_lowerCAmelCase								=lm_model(input_ids=__UpperCAmelCase  )
				_lowerCAmelCase								=(batch_size, input_ids.shape[1], config.vocab_size)
				self.assertEqual(outputs["""logits"""].shape     ,     __UpperCAmelCase  )



			def 					_lowerCAmelCase		(							self  ) ->     int:
				_lowerCAmelCase								=MBartConfig(
				    vocab_size=self.vocab_size     ,     d_model=14     ,     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=48     ,     )
				_lowerCAmelCase								=FlaxMBartForConditionalGeneration(__UpperCAmelCase  )
				_lowerCAmelCase								=np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]     ,     dtype=np.intaa  )
				_lowerCAmelCase								=np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]     ,     dtype=np.intaa  )
				_lowerCAmelCase								=lm_model(input_ids=__UpperCAmelCase     ,     decoder_input_ids=__UpperCAmelCase  )
				_lowerCAmelCase								=(*summary.shape, config.vocab_size)
				self.assertEqual(outputs["""logits"""].shape     ,     __UpperCAmelCase  )



			def 					_lowerCAmelCase		(							self  ) ->     int:
				_lowerCAmelCase								=np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]     ,     dtype=np.intaa  )
				_lowerCAmelCase								=shift_tokens_right(__UpperCAmelCase     ,     1  )
				_lowerCAmelCase								=np.equal(__UpperCAmelCase     ,     1  ).astype(np.floataa  ).sum()
				_lowerCAmelCase								=np.equal(__UpperCAmelCase     ,     1  ).astype(np.floataa  ).sum()
				self.assertEqual(shifted.shape     ,     input_ids.shape  )
				self.assertEqual(__UpperCAmelCase     ,     n_pad_before - 1  )
				self.assertTrue(np.equal(shifted[:, 0]     ,     2  ).all()  )







@require_flax
class 		lowerCamelCase__					(					__magic_name__ ,	unittest.TestCase ,	__magic_name__				):





			'''simple docstring'''
			lowerCamelCase  					=		True
			lowerCamelCase  					=		(
			    (
			        FlaxMBartModel,
			        FlaxMBartForConditionalGeneration,
			        FlaxMBartForSequenceClassification,
			        FlaxMBartForQuestionAnswering,
			    )
			    if is_flax_available()
			    else ()
			)
			lowerCamelCase  					=		(FlaxMBartForConditionalGeneration,) if is_flax_available() else ()



			def 					_lowerCAmelCase		(							self  ) ->     Dict:
				_lowerCAmelCase								=FlaxMBartModelTester(self  )



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



			def 					_lowerCAmelCase		(							self  ) ->     Dict:
				_lowerCAmelCase						,							_lowerCAmelCase								=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(__UpperCAmelCase     ,     __UpperCAmelCase     ,     __UpperCAmelCase  )



			def 					_lowerCAmelCase		(							self  ) ->     int:
				_lowerCAmelCase						,							_lowerCAmelCase								=self.model_tester.prepare_config_and_inputs_for_common()

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

						@jax.jit
						def encode_jitted(__UpperCAmelCase     ,     __UpperCAmelCase=None     ,     **__UpperCAmelCase  ):
							return model.encode(input_ids=__UpperCAmelCase     ,     attention_mask=__UpperCAmelCase  )

						with self.subTest("""JIT Enabled"""  ):
							_lowerCAmelCase								=encode_jitted(**__UpperCAmelCase  ).to_tuple()

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

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



			def 					_lowerCAmelCase		(							self  ) ->     Optional[Any]:
				_lowerCAmelCase						,							_lowerCAmelCase								=self.model_tester.prepare_config_and_inputs_for_common()

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

						_lowerCAmelCase								={
						    """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(__UpperCAmelCase     ,     __UpperCAmelCase     ,     __UpperCAmelCase  ):
							return model.decode(
							    decoder_input_ids=__UpperCAmelCase     ,     decoder_attention_mask=__UpperCAmelCase     ,     encoder_outputs=__UpperCAmelCase     ,     )

						with self.subTest("""JIT Enabled"""  ):
							_lowerCAmelCase								=decode_jitted(**__UpperCAmelCase  ).to_tuple()

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

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



			@slow
			def 					_lowerCAmelCase		(							self  ) ->     Tuple:
				for model_class_name in self.all_model_classes:
					_lowerCAmelCase								=model_class_name.from_pretrained("""facebook/mbart-large-cc25"""     ,     from_pt=__UpperCAmelCase  )
					# FlaxMBartForSequenceClassification expects eos token in input_ids
					_lowerCAmelCase								=np.ones((1, 1)  ) * model.config.eos_token_id
					_lowerCAmelCase								=model(__UpperCAmelCase  )
					self.assertIsNotNone(__UpperCAmelCase  )







@require_flax
@require_sentencepiece
@require_tokenizers
class 		lowerCamelCase__					(					unittest.TestCase				):





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



			@cached_property
			def 					_lowerCAmelCase		(							self  ) ->     List[str]:
				return AutoTokenizer.from_pretrained(self.model_name  )



			@cached_property
			def 					_lowerCAmelCase		(							self  ) ->     int:
				_lowerCAmelCase								=FlaxMBartForConditionalGeneration.from_pretrained(self.model_name     ,     from_pt=__UpperCAmelCase  )
				return model



			def 					_lowerCAmelCase		(							self     ,     **__UpperCAmelCase  ) ->     str:
				_lowerCAmelCase								=self.translate_src_text(**__UpperCAmelCase  )
				self.assertListEqual(self.expected_text     ,     __UpperCAmelCase  )



			def 					_lowerCAmelCase		(							self     ,     **__UpperCAmelCase  ) ->     str:
				_lowerCAmelCase								=self.tokenizer(self.src_text     ,     **__UpperCAmelCase     ,     return_tensors="""np"""  )
				_lowerCAmelCase								=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=__UpperCAmelCase     ,     num_beams=2     ,     ).sequences
				_lowerCAmelCase								=self.tokenizer.batch_decode(__UpperCAmelCase     ,     skip_special_tokens=__UpperCAmelCase  )
				return generated_words



			@slow
			def 					_lowerCAmelCase		(							self  ) ->     List[Any]:
				self._assert_generated_batch_equal_expected()





"""simple docstring"""


import itertools



def     _lowerCamelCase(__UpperCamelCase       )				-> List[str]:
	return (
	    int("""""".join(combination[0:2]       )       ) * int("""""".join(combination[2:5]       )       )
	    == int("""""".join(combination[5:9]       )       )
	) or (
	    int("""""".join(combination[0]       )       ) * int("""""".join(combination[1:5]       )       )
	    == int("""""".join(combination[5:9]       )       )
	)



def     _lowerCamelCase()				-> Dict:
	return sum(
	    {
	        int("""""".join(pandigital[5:9]       )       )
	        for pandigital in itertools.permutations("""123456789"""       )
	        if is_combination_valid(__UpperCamelCase       )
	    }       )


if __name__ == "__main__":
			print(solution())



