
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_torch_available


_UpperCAmelCase   :					Optional[Any]    =			{'configuration_bert_generation': ['BertGenerationConfig']}

try:
					if not is_sentencepiece_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					_UpperCAmelCase   :					List[Any]    =			['BertGenerationTokenizer']

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					_UpperCAmelCase   :					int    =			[
					    'BertGenerationDecoder',
					    'BertGenerationEncoder',
					    'BertGenerationPreTrainedModel',
					    'load_tf_weights_in_bert_generation',
					]


if TYPE_CHECKING:
					from .configuration_bert_generation import BertGenerationConfig

					try:
										if not is_sentencepiece_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_bert_generation import BertGenerationTokenizer

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_bert_generation import (
										    BertGenerationDecoder,
										    BertGenerationEncoder,
										    BertGenerationPreTrainedModel,
										    load_tf_weights_in_bert_generation,
										)

else:
					import sys

					_UpperCAmelCase   :					Optional[int]    =			_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)







'''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 :			Dict							   =      '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    lowercase_   (		lowercase__					,	lowercase__					,	lowercase__=None					,	lowercase__=None					,	lowercase__=None					,	lowercase__=None					,	lowercase__=None					,	lowercase__=None					,	) ->List[Any]:
						if attention_mask is None:
												_snake_case: Optional[int]	       =		np.where(input_ids != config.pad_token_id					,	1					,	0    )
						if decoder_attention_mask is None:
												_snake_case: Optional[Any]	       =		np.where(decoder_input_ids != config.pad_token_id					,	1					,	0    )
						if head_mask is None:
												_snake_case: Union[str, Any]	       =		np.ones((config.encoder_layers, config.encoder_attention_heads)    )
						if decoder_head_mask is None:
												_snake_case: List[str]	       =		np.ones((config.decoder_layers, config.decoder_attention_heads)    )
						if cross_attn_head_mask is None:
												_snake_case: Union[str, Any]	       =		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	):


					def __init__(       self						:	Optional[Any]							,  __snake_case						:	Optional[Any]							,  __snake_case						:	List[Any]=13							,  __snake_case						:	Dict=7							,  __snake_case						:	int=True							,  __snake_case						:	Tuple=False							,  __snake_case						:	Optional[int]=99							,  __snake_case						:	Optional[Any]=16							,  __snake_case						:	Any=2							,  __snake_case						:	Dict=4							,  __snake_case						:	int=4							,  __snake_case						:	Any="gelu"							,  __snake_case						:	List[Any]=0.1							,  __snake_case						:	List[str]=0.1							,  __snake_case						:	int=32							,  __snake_case						:	Optional[Any]=2							,  __snake_case						:	int=1							,  __snake_case						:	str=0							,  __snake_case						:	Any=2							,  __snake_case						:	Optional[Any]=0.02							,  ):





											'''simple docstring'''





											_snake_case: Dict	       =		parent
											_snake_case: List[Any]	       =		batch_size
											_snake_case: Union[str, Any]	       =		seq_length
											_snake_case: List[str]	       =		is_training
											_snake_case: Optional[int]	       =		use_labels
											_snake_case: List[str]	       =		vocab_size
											_snake_case: Union[str, Any]	       =		hidden_size
											_snake_case: Optional[Any]	       =		num_hidden_layers
											_snake_case: int	       =		num_attention_heads
											_snake_case: Tuple	       =		intermediate_size
											_snake_case: Optional[int]	       =		hidden_act
											_snake_case: Optional[int]	       =		hidden_dropout_prob
											_snake_case: Union[str, Any]	       =		attention_probs_dropout_prob
											_snake_case: Tuple	       =		max_position_embeddings
											_snake_case: List[Any]	       =		eos_token_id
											_snake_case: int	       =		pad_token_id
											_snake_case: Optional[Any]	       =		bos_token_id
											_snake_case: Any	       =		decoder_start_token_id
											_snake_case: Dict	       =		initializer_range


					def  SCREAMING_SNAKE_CASE_		(       self						:	Dict							):





											'''simple docstring'''





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

											_snake_case: List[Any]	       =		shift_tokens_right(__snake_case							,  1							)

											_snake_case: Dict	       =		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=__snake_case							,  )
											_snake_case: Optional[Any]	       =		prepare_mbart_inputs_dict(__snake_case							,  __snake_case							,  __snake_case							)
											return config, inputs_dict


					def  SCREAMING_SNAKE_CASE_		(       self						:	Any							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: int	       =		self.prepare_config_and_inputs()
											return config, inputs_dict


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[str]							,  __snake_case						:	str							,  __snake_case						:	List[Any]							,  __snake_case						:	Optional[int]							):





											'''simple docstring'''





											_snake_case: List[str]	       =		20
											_snake_case: Optional[Any]	       =		model_class_name(__snake_case							)

											_snake_case: List[Any]	       =		model.encode(inputs_dict['input_ids']							)

											_snake_case       ,    _snake_case: str	       =		(
											    inputs_dict['decoder_input_ids'],
											    inputs_dict['decoder_attention_mask'],
											)

											_snake_case: str	       =		model.init_cache(decoder_input_ids.shape[0]							,  __snake_case							,  __snake_case							)
											_snake_case: Tuple	       =		jnp.ones((decoder_input_ids.shape[0], max_decoder_length)							,  dtype='i4'							)

											_snake_case: Union[str, Any]	       =		jnp.broadcast_to(
											    jnp.arange(decoder_input_ids.shape[-1] - 1							)[None, :]							,  (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)							,  )
											_snake_case: Dict	       =		model.decode(
											    decoder_input_ids[:, :-1]							,  __snake_case							,  decoder_attention_mask=__snake_case							,  past_key_values=__snake_case							,  decoder_position_ids=__snake_case							,  )

											_snake_case: Optional[int]	       =		jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]							,  dtype='i4'							)
											_snake_case: List[str]	       =		model.decode(
											    decoder_input_ids[:, -1:]							,  __snake_case							,  decoder_attention_mask=__snake_case							,  past_key_values=outputs_cache.past_key_values							,  decoder_position_ids=__snake_case							,  )

											_snake_case: Any	       =		model.decode(__snake_case							,  __snake_case							)

											_snake_case: List[Any]	       =		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  SCREAMING_SNAKE_CASE_		(       self						:	Optional[int]							,  __snake_case						:	int							,  __snake_case						:	int							,  __snake_case						:	int							):





											'''simple docstring'''





											_snake_case: Dict	       =		20
											_snake_case: Any	       =		model_class_name(__snake_case							)

											_snake_case: Dict	       =		model.encode(inputs_dict['input_ids']							)

											_snake_case       ,    _snake_case: Tuple	       =		(
											    inputs_dict['decoder_input_ids'],
											    inputs_dict['decoder_attention_mask'],
											)

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

											_snake_case: str	       =		model.init_cache(decoder_input_ids.shape[0]							,  __snake_case							,  __snake_case							)
											_snake_case: Dict	       =		jnp.broadcast_to(
											    jnp.arange(decoder_input_ids.shape[-1] - 1							)[None, :]							,  (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)							,  )

											_snake_case: int	       =		model.decode(
											    decoder_input_ids[:, :-1]							,  __snake_case							,  decoder_attention_mask=__snake_case							,  past_key_values=__snake_case							,  decoder_position_ids=__snake_case							,  )
											_snake_case: Any	       =		jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]							,  dtype='i4'							)
											_snake_case: str	       =		model.decode(
											    decoder_input_ids[:, -1:]							,  __snake_case							,  past_key_values=outputs_cache.past_key_values							,  decoder_attention_mask=__snake_case							,  decoder_position_ids=__snake_case							,  )

											_snake_case: Any	       =		model.decode(__snake_case							,  __snake_case							,  decoder_attention_mask=__snake_case							)

											_snake_case: Optional[int]	       =		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	):
					_SCREAMING_SNAKE_CASE				 =     99


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[str]							):





											'''simple docstring'''





											_snake_case: int	       =		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							,  )

											_snake_case: Any	       =		input_ids.shape[0]
											_snake_case: str	       =		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  SCREAMING_SNAKE_CASE_		(       self						:	Union[str, Any]							):





											'''simple docstring'''





											_snake_case       ,    _snake_case       ,    _snake_case: Tuple	       =		self._get_config_and_data()
											_snake_case: Any	       =		FlaxMBartForSequenceClassification(__snake_case							)
											_snake_case: List[Any]	       =		model(input_ids=__snake_case							,  decoder_input_ids=__snake_case							)
											_snake_case: List[str]	       =		(batch_size, config.num_labels)
											self.assertEqual(outputs['logits'].shape							,  __snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[Any]							):





											'''simple docstring'''





											_snake_case       ,    _snake_case       ,    _snake_case: Union[str, Any]	       =		self._get_config_and_data()
											_snake_case: List[str]	       =		FlaxMBartForQuestionAnswering(__snake_case							)
											_snake_case: Union[str, Any]	       =		model(input_ids=__snake_case							)

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


					def  SCREAMING_SNAKE_CASE_		(       self						:	Tuple							):





											'''simple docstring'''





											_snake_case       ,    _snake_case       ,    _snake_case: Any	       =		self._get_config_and_data()
											_snake_case: Optional[int]	       =		FlaxMBartForConditionalGeneration(__snake_case							)
											_snake_case: List[str]	       =		lm_model(input_ids=__snake_case							)
											_snake_case: List[Any]	       =		(batch_size, input_ids.shape[1], config.vocab_size)
											self.assertEqual(outputs['logits'].shape							,  __snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[Any]							):





											'''simple docstring'''





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




					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[int]							):





											'''simple docstring'''





											_snake_case: Any	       =		np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]							,  dtype=np.intaa							)
											_snake_case: Any	       =		shift_tokens_right(__snake_case							,  1							)
											_snake_case: List[str]	       =		np.equal(__snake_case							,  1							).astype(np.floataa							).sum()
											_snake_case: List[str]	       =		np.equal(__snake_case							,  1							).astype(np.floataa							).sum()
											self.assertEqual(shifted.shape							,  input_ids.shape							)
											self.assertEqual(__snake_case							,  n_pad_before - 1							)
											self.assertTrue(np.equal(shifted[:, 0]							,  2							).all()							)







@require_flax
class       lowerCamelCase       (    __UpperCAmelCase  ,				unittest.TestCase  ,				__UpperCAmelCase	):
					_SCREAMING_SNAKE_CASE				 =     True
					_SCREAMING_SNAKE_CASE				 =     (
					    (
					        FlaxMBartModel,
					        FlaxMBartForConditionalGeneration,
					        FlaxMBartForSequenceClassification,
					        FlaxMBartForQuestionAnswering,
					    )
					    if is_flax_available()
					    else ()
					)
					_SCREAMING_SNAKE_CASE				 =     (FlaxMBartForConditionalGeneration,) if is_flax_available() else ()


					def  SCREAMING_SNAKE_CASE_		(       self						:	Tuple							):





											'''simple docstring'''





											_snake_case: Union[str, Any]	       =		FlaxMBartModelTester(self							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	int							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: int	       =		self.model_tester.prepare_config_and_inputs()
											for model_class in self.all_model_classes:
																	self.model_tester.check_use_cache_forward(__snake_case							,  __snake_case							,  __snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	Union[str, Any]							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: Tuple	       =		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(__snake_case							,  __snake_case							,  __snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[str]							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: Any	       =		self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	with self.subTest(model_class.__name__							):
																							_snake_case: Dict	       =		self._prepare_for_class(__snake_case							,  __snake_case							)
																							_snake_case: List[Any]	       =		model_class(__snake_case							)

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

																							with self.subTest('JIT Enabled'							):
																													_snake_case: List[str]	       =		encode_jitted(**__snake_case							).to_tuple()

																							with self.subTest('JIT Disabled'							):
																													with jax.disable_jit():
																																			_snake_case: List[Any]	       =		encode_jitted(**__snake_case							).to_tuple()

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


					def  SCREAMING_SNAKE_CASE_		(       self						:	str							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: Optional[int]	       =		self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	with self.subTest(model_class.__name__							):
																							_snake_case: Tuple	       =		model_class(__snake_case							)
																							_snake_case: Dict	       =		model.encode(inputs_dict['input_ids']							,  inputs_dict['attention_mask']							)

																							_snake_case: Dict	       =		{
																							    '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(__snake_case						:	Dict							,  __snake_case						:	Dict							,  __snake_case						:	Union[str, Any]							):
																													return model.decode(
																													    decoder_input_ids=__snake_case							,  decoder_attention_mask=__snake_case							,  encoder_outputs=__snake_case							,  )

																							with self.subTest('JIT Enabled'							):
																													_snake_case: str	       =		decode_jitted(**__snake_case							).to_tuple()

																							with self.subTest('JIT Disabled'							):
																													with jax.disable_jit():
																																			_snake_case: Dict	       =		decode_jitted(**__snake_case							).to_tuple()

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




					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	int							):





											'''simple docstring'''





											for model_class_name in self.all_model_classes:
																	_snake_case: List[str]	       =		model_class_name.from_pretrained('facebook/mbart-large-cc25'							,  from_pt=__snake_case							)
																	# FlaxMBartForSequenceClassification expects eos token in input_ids
																	_snake_case: List[str]	       =		np.ones((1, 1)							) * model.config.eos_token_id
																	_snake_case: Optional[Any]	       =		model(__snake_case							)
																	self.assertIsNotNone(__snake_case							)







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


					@cached_property
					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[int]							):





											'''simple docstring'''





											return AutoTokenizer.from_pretrained(self.model_name							)


					@cached_property
					def  SCREAMING_SNAKE_CASE_		(       self						:	Dict							):





											'''simple docstring'''





											_snake_case: Dict	       =		FlaxMBartForConditionalGeneration.from_pretrained(self.model_name							,  from_pt=__snake_case							)
											return model


					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[int]							,  **__snake_case						:	int							):





											'''simple docstring'''





											_snake_case: List[Any]	       =		self.translate_src_text(**__snake_case							)
											self.assertListEqual(self.expected_text							,  __snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	str							,  **__snake_case						:	Tuple							):





											'''simple docstring'''





											_snake_case: int	       =		self.tokenizer(self.src_text							,  **__snake_case							,  return_tensors='np'							)
											_snake_case: Optional[Any]	       =		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=__snake_case							,  num_beams=2							,  ).sequences
											_snake_case: Dict	       =		self.tokenizer.batch_decode(__snake_case							,  skip_special_tokens=__snake_case							)
											return generated_words




					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[Any]							):





											'''simple docstring'''





											self._assert_generated_batch_equal_expected()





