




'''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
				_snake_case      :      str   							=		'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 		snake_case_    (UpperCamelCase       :  Optional[int] ,						UpperCamelCase       :  int ,						UpperCamelCase       :  Any=None ,						UpperCamelCase       :  Any=None ,						UpperCamelCase       :  List[str]=None ,						UpperCamelCase       :  Union[str, Any]=None ,						UpperCamelCase       :  Union[str, Any]=None ,						UpperCamelCase       :  Tuple=None ,						):





							'''simple docstring'''


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







							def __init__(  self       :  List[Any]							,			lowerCAmelCase_       :  Optional[Any]							,			lowerCAmelCase_       :  Union[str, Any]=13							,			lowerCAmelCase_       :  Union[str, Any]=7							,			lowerCAmelCase_       :  Tuple=True							,			lowerCAmelCase_       :  int=False							,			lowerCAmelCase_       :  Dict=99							,			lowerCAmelCase_       :  int=16							,			lowerCAmelCase_       :  Optional[Any]=2							,			lowerCAmelCase_       :  Any=4							,			lowerCAmelCase_       :  str=4							,			lowerCAmelCase_       :  Dict="gelu"							,			lowerCAmelCase_       :  Tuple=0.1							,			lowerCAmelCase_       :  Dict=0.1							,			lowerCAmelCase_       :  Any=32							,			lowerCAmelCase_       :  Dict=2							,			lowerCAmelCase_       :  Union[str, Any]=1							,			lowerCAmelCase_       :  Tuple=0							,			lowerCAmelCase_       :  Tuple=2							,			lowerCAmelCase_       :  str=0.0_2							,			)			->							Tuple:


														"""simple docstring"""


														_a	     =				parent
														_a	     =				batch_size
														_a	     =				seq_length
														_a	     =				is_training
														_a	     =				use_labels
														_a	     =				vocab_size
														_a	     =				hidden_size
														_a	     =				num_hidden_layers
														_a	     =				num_attention_heads
														_a	     =				intermediate_size
														_a	     =				hidden_act
														_a	     =				hidden_dropout_prob
														_a	     =				attention_probs_dropout_prob
														_a	     =				max_position_embeddings
														_a	     =				eos_token_id
														_a	     =				pad_token_id
														_a	     =				bos_token_id
														_a	     =				decoder_start_token_id
														_a	     =				initializer_range







							def    __lowerCAmelCase			(  self       :  Optional[int]					)			->							Any:


														"""simple docstring"""


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

														_a	     =				shift_tokens_right(__UpperCamelCase							,			1					)

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







							def    __lowerCAmelCase			(  self       :  Tuple					)			->							List[Any]:


														"""simple docstring"""


														_a     ,							_a	     =				self.prepare_config_and_inputs()
														return config, inputs_dict







							def    __lowerCAmelCase			(  self       :  int							,			lowerCAmelCase_       :  Optional[int]							,			lowerCAmelCase_       :  Union[str, Any]							,			lowerCAmelCase_       :  List[str]					)			->							str:


														"""simple docstring"""


														_a	     =				20
														_a	     =				model_class_name(__UpperCamelCase					)

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

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

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

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

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

														_a	     =				model.decode(__UpperCamelCase							,			__UpperCamelCase					)

														_a	     =				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       :  Optional[int]							,			lowerCAmelCase_       :  int							,			lowerCAmelCase_       :  Dict							,			lowerCAmelCase_       :  Union[str, Any]					)			->							Optional[int]:


														"""simple docstring"""


														_a	     =				20
														_a	     =				model_class_name(__UpperCamelCase					)

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

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

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

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

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

														_a	     =				model.decode(__UpperCamelCase							,			__UpperCamelCase							,			decoder_attention_mask=__UpperCamelCase					)

														_a	     =				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 			A    (		unittest.TestCase    ):
							lowercase_							=				99







							def    __lowerCAmelCase			(  self       :  Union[str, Any]					)			->							Dict:


														"""simple docstring"""


														_a	     =				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							,			)

														_a	     =				input_ids.shape[0]
														_a	     =				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       :  Tuple					)			->							Optional[int]:


														"""simple docstring"""


														_a     ,							_a     ,							_a	     =				self._get_config_and_data()
														_a	     =				FlaxMBartForSequenceClassification(__UpperCamelCase					)
														_a	     =				model(input_ids=__UpperCamelCase							,			decoder_input_ids=__UpperCamelCase					)
														_a	     =				(batch_size, config.num_labels)
														self.assertEqual(outputs['''logits'''].shape							,			__UpperCamelCase					)







							def    __lowerCAmelCase			(  self       :  int					)			->							Optional[int]:


														"""simple docstring"""


														_a     ,							_a     ,							_a	     =				self._get_config_and_data()
														_a	     =				FlaxMBartForQuestionAnswering(__UpperCamelCase					)
														_a	     =				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       :  List[Any]					)			->							str:


														"""simple docstring"""


														_a     ,							_a     ,							_a	     =				self._get_config_and_data()
														_a	     =				FlaxMBartForConditionalGeneration(__UpperCamelCase					)
														_a	     =				lm_model(input_ids=__UpperCamelCase					)
														_a	     =				(batch_size, input_ids.shape[1], config.vocab_size)
														self.assertEqual(outputs['''logits'''].shape							,			__UpperCamelCase					)







							def    __lowerCAmelCase			(  self       :  int					)			->							Dict:


														"""simple docstring"""


														_a	     =				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							,			)
														_a	     =				FlaxMBartForConditionalGeneration(__UpperCamelCase					)
														_a	     =				np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]							,			dtype=np.intaa					)
														_a	     =				np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]							,			dtype=np.intaa					)
														_a	     =				lm_model(input_ids=__UpperCamelCase							,			decoder_input_ids=__UpperCamelCase					)
														_a	     =				(*summary.shape, config.vocab_size)
														self.assertEqual(outputs['''logits'''].shape							,			__UpperCamelCase					)







							def    __lowerCAmelCase			(  self       :  Any					)			->							List[str]:


														"""simple docstring"""


														_a	     =				np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]							,			dtype=np.intaa					)
														_a	     =				shift_tokens_right(__UpperCamelCase							,			1					)
														_a	     =				np.equal(__UpperCamelCase							,			1					).astype(np.floataa					).sum()
														_a	     =				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 			A    (		UpperCAmelCase_     ,unittest.TestCase     ,UpperCAmelCase_    ):
							lowercase_							=				True
							lowercase_							=				(
							    (
							        FlaxMBartModel,
							        FlaxMBartForConditionalGeneration,
							        FlaxMBartForSequenceClassification,
							        FlaxMBartForQuestionAnswering,
							    )
							    if is_flax_available()
							    else ()
							)
							lowercase_							=				(FlaxMBartForConditionalGeneration,) if is_flax_available() else ()







							def    __lowerCAmelCase			(  self       :  Dict					)			->							Tuple:


														"""simple docstring"""


														_a	     =				FlaxMBartModelTester(self					)







							def    __lowerCAmelCase			(  self       :  Union[str, Any]					)			->							Optional[Any]:


														"""simple docstring"""


														_a     ,							_a	     =				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       :  Optional[int]					)			->							List[str]:


														"""simple docstring"""


														_a     ,							_a	     =				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       :  List[Any]					)			->							Dict:


														"""simple docstring"""


														_a     ,							_a	     =				self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					with self.subTest(model_class.__name__					):
																												_a	     =				self._prepare_for_class(__UpperCamelCase							,			__UpperCamelCase					)
																												_a	     =				model_class(__UpperCamelCase					)

																												@jax.jit
																												def encode_jitted(lowerCAmelCase_       :  Optional[int]							,			lowerCAmelCase_       :  Any=None							,			**lowerCAmelCase_       :  str					):
																																			return model.encode(input_ids=__UpperCamelCase							,			attention_mask=__UpperCamelCase					)

																												with self.subTest('''JIT Enabled'''					):
																																			_a	     =				encode_jitted(**__UpperCamelCase					).to_tuple()

																												with self.subTest('''JIT Disabled'''					):
																																			with jax.disable_jit():
																																										_a	     =				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       :  Any					)			->							Any:


														"""simple docstring"""


														_a     ,							_a	     =				self.model_tester.prepare_config_and_inputs_for_common()

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

																												_a	     =				{
																												    '''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_       :  Optional[int]							,			lowerCAmelCase_       :  List[Any]							,			lowerCAmelCase_       :  List[str]					):
																																			return model.decode(
																																			    decoder_input_ids=__UpperCamelCase							,			decoder_attention_mask=__UpperCamelCase							,			encoder_outputs=__UpperCamelCase							,			)

																												with self.subTest('''JIT Enabled'''					):
																																			_a	     =				decode_jitted(**__UpperCamelCase					).to_tuple()

																												with self.subTest('''JIT Disabled'''					):
																																			with jax.disable_jit():
																																										_a	     =				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       :  List[Any]					)			->							List[str]:


														"""simple docstring"""


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



@require_flax
@require_sentencepiece
@require_tokenizers
class 			A    (		unittest.TestCase    ):
							lowercase_							=				[
							    ' UN Chief Says There Is No Military Solution in Syria',
							]
							lowercase_							=				[
							    'Şeful ONU declară că nu există o soluţie militară în Siria',
							]
							lowercase_							=				'facebook/mbart-large-en-ro'







							@cached_property
							def    __lowerCAmelCase			(  self       :  Dict					)			->							Tuple:


														"""simple docstring"""


														return AutoTokenizer.from_pretrained(self.model_name					)







							@cached_property
							def    __lowerCAmelCase			(  self       :  Optional[int]					)			->							Dict:


														"""simple docstring"""


														_a	     =				FlaxMBartForConditionalGeneration.from_pretrained(self.model_name							,			from_pt=__UpperCamelCase					)
														return model







							def    __lowerCAmelCase			(  self       :  Optional[int]							,			**lowerCAmelCase_       :  Any					)			->							Dict:


														"""simple docstring"""


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







							def    __lowerCAmelCase			(  self       :  Union[str, Any]							,			**lowerCAmelCase_       :  Optional[Any]					)			->							List[Any]:


														"""simple docstring"""


														_a	     =				self.tokenizer(self.src_text							,			**__UpperCamelCase							,			return_tensors='''np'''					)
														_a	     =				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
														_a	     =				self.tokenizer.batch_decode(__UpperCamelCase							,			skip_special_tokens=__UpperCamelCase					)
														return generated_words







							@slow
							def    __lowerCAmelCase			(  self       :  Any					)			->							Union[str, Any]:


														"""simple docstring"""


														self._assert_generated_batch_equal_expected()







'''simple docstring'''







from functools import lru_cache






@lru_cache
def 		_lowercase		(	lowercase_  ):



		'''simple docstring'''







		if not isinstance(lowercase_						,							lowercase_  ):
				raise TypeError("k must be an integer."  )
		if k < 0:
				raise ValueError("k integer must be greater or equal to zero."  )
		if k == 0:
				return 0
		elif k == 1:
				return 1
		else:
				return fibonacci(k - 1  ) + fibonacci(k - 2  )






def 		_lowercase		(	lowercase_						,							lowercase_  ):



		'''simple docstring'''







		_snake_case	       =       len(lowercase_  )
		# Find m such that F_m >= n where F_i is the i_th fibonacci number.
		_snake_case	       =       0
		while True:
				if fibonacci(lowercase_  ) >= len_list:
						_snake_case	       =       i
						break
				i += 1
		_snake_case	       =       0
		while fibb_k > 0:
				_snake_case	       =       min(
				    offset + fibonacci(fibb_k - 1  )						,							len_list - 1  )  # Prevent out of range
				_snake_case	       =       arr[index_k]
				if item_k_1 == val:
						return index_k
				elif val < item_k_1:
						fibb_k -= 1
				elif val > item_k_1:
						offset += fibonacci(fibb_k - 1  )
						fibb_k -= 2
		else:
				return -1


if __name__ == "__main__":
						import doctest

						doctest.testmod()





