
"""simple docstring"""

import inspect
import unittest

from transformers import EfficientNetConfig
from transformers.testing_utils import is_pipeline_test, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import EfficientNetForImageClassification, EfficientNetModel
					from transformers.models.efficientnet.modeling_efficientnet import EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					from PIL import Image

					from transformers import AutoImageProcessor






class        _lowercase   :




	def __init__( self		, UpperCAmelCase_		, UpperCAmelCase_=13		, UpperCAmelCase_=32		, UpperCAmelCase_=3		, UpperCAmelCase_=[3, 3, 5]		, UpperCAmelCase_=[32, 16, 24]		, UpperCAmelCase_=[16, 24, 40]		, UpperCAmelCase_=[1, 1, 2]		, UpperCAmelCase_=[1, 1, 2]		, UpperCAmelCase_=[1, 6, 6]		, UpperCAmelCase_=True		, UpperCAmelCase_=True		, UpperCAmelCase_=37		, UpperCAmelCase_="gelu"		, UpperCAmelCase_=10		, )		->							Union[str, Any]:
			lowerCamelCase     :			Optional[int]		=    parent
			lowerCamelCase     :			int		=    batch_size
			lowerCamelCase     :			Tuple		=    image_size
			lowerCamelCase     :			List[Any]		=    num_channels
			lowerCamelCase     :			str		=    kernel_sizes
			lowerCamelCase     :			Dict		=    in_channels
			lowerCamelCase     :			Tuple		=    out_channels
			lowerCamelCase     :			Any		=    strides
			lowerCamelCase     :			List[Any]		=    num_block_repeats
			lowerCamelCase     :			int		=    expand_ratios
			lowerCamelCase     :			List[Any]		=    is_training
			lowerCamelCase     :			int		=    hidden_act
			lowerCamelCase     :			Union[str, Any]		=    num_labels
			lowerCamelCase     :			Optional[Any]		=    use_labels




	def 			_UpperCamelCase    ( self					)		->							Dict:
			lowerCamelCase     :			Tuple		=    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

			lowerCamelCase     :			Union[str, Any]		=    None
			if self.use_labels:
					lowerCamelCase     :			Union[str, Any]		=    ids_tensor([self.batch_size]		, self.num_labels					)

			lowerCamelCase     :			Optional[Any]		=    self.get_config()
			return config, pixel_values, labels




	def 			_UpperCamelCase    ( self					)		->							Optional[Any]:
			return EfficientNetConfig(
			    num_channels=self.num_channels		, kernel_sizes=self.kernel_sizes		, in_channels=self.in_channels		, out_channels=self.out_channels		, strides=self.strides		, num_block_repeats=self.num_block_repeats		, expand_ratios=self.expand_ratios		, hidden_act=self.hidden_act		, num_labels=self.num_labels		, )




	def 			_UpperCamelCase    ( self		, UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_					)		->							int:
			lowerCamelCase     :			List[Any]		=    EfficientNetModel(config=UpperCAmelCase_					)
			model.to(UpperCAmelCase_					)
			model.eval()
			lowerCamelCase     :			Optional[int]		=    model(UpperCAmelCase_					)
			# expected last hidden states: B, C, H // 4, W // 4
			self.parent.assertEqual(
			    result.last_hidden_state.shape		, (self.batch_size, config.hidden_dim, self.image_size // 4, self.image_size // 4)		, )




	def 			_UpperCamelCase    ( self		, UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_					)		->							Any:
			lowerCamelCase     :			Dict		=    EfficientNetForImageClassification(UpperCAmelCase_					)
			model.to(UpperCAmelCase_					)
			model.eval()
			lowerCamelCase     :			Optional[int]		=    model(UpperCAmelCase_		, labels=UpperCAmelCase_					)
			self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.num_labels)					)




	def 			_UpperCamelCase    ( self					)		->							List[str]:
			lowerCamelCase     :			Optional[int]		=    self.prepare_config_and_inputs()
			lowerCamelCase	,    lowerCamelCase	,    lowerCamelCase     :			Dict		=    config_and_inputs
			lowerCamelCase     :			Optional[Any]		=    {'pixel_values': pixel_values}
			return config, inputs_dict






@require_torch
class        _lowercase   (				__UpperCAmelCase				,       __UpperCAmelCase				,       unittest.TestCase		):
	lowercase_        =	(EfficientNetModel, EfficientNetForImageClassification) if is_torch_available() else ()
	lowercase_        =	(
	    {'feature-extraction': EfficientNetModel, 'image-classification': EfficientNetForImageClassification}
	    if is_torch_available()
	    else {}
	)

	lowercase_        =	False
	lowercase_        =	False
	lowercase_        =	False
	lowercase_        =	False
	lowercase_        =	False




	def 			_UpperCamelCase    ( self					)		->							Union[str, Any]:
			lowerCamelCase     :			int		=    EfficientNetModelTester(self					)
			lowerCamelCase     :			str		=    ConfigTester(
			    self		, config_class=UpperCAmelCase_		, has_text_modality=UpperCAmelCase_		, hidden_size=37					)




	def 			_UpperCamelCase    ( self					)		->							Dict:
			self.create_and_test_config_common_properties()
			self.config_tester.create_and_test_config_to_json_string()
			self.config_tester.create_and_test_config_to_json_file()
			self.config_tester.create_and_test_config_from_and_save_pretrained()
			self.config_tester.create_and_test_config_with_num_labels()
			self.config_tester.check_config_can_be_init_without_params()
			self.config_tester.check_config_arguments_init()




	def 			_UpperCamelCase    ( self					)		->							Tuple:
			return




	@unittest.skip(reason='EfficientNet does not use inputs_embeds'					)
	def 			_UpperCamelCase    ( self					)		->							Dict:
			pass




	@unittest.skip(reason='EfficientNet does not support input and output embeddings'					)
	def 			_UpperCamelCase    ( self					)		->							List[Any]:
			pass




	@unittest.skip(reason='EfficientNet does not use feedforward chunking'					)
	def 			_UpperCamelCase    ( self					)		->							str:
			pass




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

			for model_class in self.all_model_classes:
					lowerCamelCase     :			Union[str, Any]		=    model_class(UpperCAmelCase_					)
					lowerCamelCase     :			List[str]		=    inspect.signature(model.forward					)
					# signature.parameters is an OrderedDict => so arg_names order is deterministic
					lowerCamelCase     :			int		=    [*signature.parameters.keys()]

					lowerCamelCase     :			List[Any]		=    ['pixel_values']
					self.assertListEqual(arg_names[:1]		, UpperCAmelCase_					)




	def 			_UpperCamelCase    ( self					)		->							Dict:
			lowerCamelCase     :			Optional[Any]		=    self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*UpperCAmelCase_					)




	def 			_UpperCamelCase    ( self					)		->							Dict:
			def check_hidden_states_output(UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_					):
					lowerCamelCase     :			Dict		=    model_class(UpperCAmelCase_					)
					model.to(UpperCAmelCase_					)
					model.eval()

					with torch.no_grad():
							lowerCamelCase     :			List[str]		=    model(**self._prepare_for_class(UpperCAmelCase_		, UpperCAmelCase_					)					)

					lowerCamelCase     :			Optional[Any]		=    outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states
					lowerCamelCase     :			Optional[Any]		=    sum(config.num_block_repeats					) * 4
					self.assertEqual(len(UpperCAmelCase_					)		, UpperCAmelCase_					)

					# EfficientNet's feature maps are of shape (batch_size, num_channels, height, width)
					self.assertListEqual(
					    list(hidden_states[0].shape[-2:]					)		, [self.model_tester.image_size // 2, self.model_tester.image_size // 2]		, )

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

			for model_class in self.all_model_classes:
					lowerCamelCase     :			Optional[Any]		=    True
					check_hidden_states_output(UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_					)

					# check that output_hidden_states also work using config
					del inputs_dict["output_hidden_states"]
					lowerCamelCase     :			List[str]		=    True

					check_hidden_states_output(UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_					)




	def 			_UpperCamelCase    ( self					)		->							str:
			lowerCamelCase     :			Any		=    self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase_					)




	@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'					)
	def 			_UpperCamelCase    ( self					)		->							List[Any]:
			pass




	@slow
	def 			_UpperCamelCase    ( self					)		->							str:
			for model_name in EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCamelCase     :			List[Any]		=    EfficientNetModel.from_pretrained(UpperCAmelCase_					)
					self.assertIsNotNone(UpperCAmelCase_					)




	@is_pipeline_test
	@require_vision
	@slow
	def 			_UpperCamelCase    ( self					)		->							Any:
			super().test_pipeline_image_classification()


def 		UpperCAmelCase					(		):





		'''simple docstring'''


		lowerCamelCase     :			List[Any]		=    Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'			)
		return image






@require_torch
@require_vision
class        _lowercase   (				unittest.TestCase		):




	@cached_property
	def 			_UpperCamelCase    ( self					)		->							str:
			return AutoImageProcessor.from_pretrained('google/efficientnet-b7'					) if is_vision_available() else None




	@slow
	def 			_UpperCamelCase    ( self					)		->							Any:
			lowerCamelCase     :			List[str]		=    EfficientNetForImageClassification.from_pretrained('google/efficientnet-b7'					).to(UpperCAmelCase_					)

			lowerCamelCase     :			Tuple		=    self.default_image_processor
			lowerCamelCase     :			Optional[Any]		=    prepare_img()
			lowerCamelCase     :			str		=    image_processor(images=UpperCAmelCase_		, return_tensors='pt'					).to(UpperCAmelCase_					)

			# forward pass
			with torch.no_grad():
					lowerCamelCase     :			Union[str, Any]		=    model(**UpperCAmelCase_					)

			# verify the logits
			lowerCamelCase     :			Tuple		=    torch.Size((1, 1000)					)
			self.assertEqual(outputs.logits.shape		, UpperCAmelCase_					)

			lowerCamelCase     :			Tuple		=    torch.tensor([0.0001, 0.0002, 0.0002]					).to(UpperCAmelCase_					)
			self.assertTrue(torch.allclose(outputs.logits[0, :3]		, UpperCAmelCase_		, atol=1E-4					)					)








"""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 		UpperCAmelCase					(		a_,						a_,						a_=None,						a_=None,						a_=None,						a_=None,						a_=None,						a_=None,						):





		'''simple docstring'''


		if attention_mask is None:
				lowerCamelCase     :			Dict		=    np.where(input_ids != config.pad_token_id,						1,						0			)
		if decoder_attention_mask is None:
				lowerCamelCase     :			Dict		=    np.where(decoder_input_ids != config.pad_token_id,						1,						0			)
		if head_mask is None:
				lowerCamelCase     :			List[Any]		=    np.ones((config.encoder_layers, config.encoder_attention_heads)			)
		if decoder_head_mask is None:
				lowerCamelCase     :			Optional[Any]		=    np.ones((config.decoder_layers, config.decoder_attention_heads)			)
		if cross_attn_head_mask is None:
				lowerCamelCase     :			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        _lowercase   (				unittest.TestCase		):




	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.02		, )		->							Union[str, Any]:
			lowerCamelCase     :			Any		=    parent
			lowerCamelCase     :			Dict		=    batch_size
			lowerCamelCase     :			str		=    seq_length
			lowerCamelCase     :			Dict		=    is_training
			lowerCamelCase     :			List[Any]		=    use_labels
			lowerCamelCase     :			Dict		=    vocab_size
			lowerCamelCase     :			List[str]		=    hidden_size
			lowerCamelCase     :			Tuple		=    num_hidden_layers
			lowerCamelCase     :			Union[str, Any]		=    num_attention_heads
			lowerCamelCase     :			int		=    intermediate_size
			lowerCamelCase     :			str		=    hidden_act
			lowerCamelCase     :			str		=    hidden_dropout_prob
			lowerCamelCase     :			str		=    attention_probs_dropout_prob
			lowerCamelCase     :			List[str]		=    max_position_embeddings
			lowerCamelCase     :			int		=    eos_token_id
			lowerCamelCase     :			Optional[int]		=    pad_token_id
			lowerCamelCase     :			Dict		=    bos_token_id
			lowerCamelCase     :			List[Any]		=    decoder_start_token_id
			lowerCamelCase     :			Dict		=    initializer_range




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

			lowerCamelCase     :			int		=    shift_tokens_right(UpperCAmelCase_		, 1					)

			lowerCamelCase     :			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=UpperCAmelCase_		, )
			lowerCamelCase     :			List[str]		=    prepare_mbart_inputs_dict(UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_					)
			return config, inputs_dict




	def 			_UpperCamelCase    ( self					)		->							Any:
			lowerCamelCase	,    lowerCamelCase     :			str		=    self.prepare_config_and_inputs()
			return config, inputs_dict




	def 			_UpperCamelCase    ( self		, UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_					)		->							Any:
			lowerCamelCase     :			List[str]		=    20
			lowerCamelCase     :			Optional[int]		=    model_class_name(UpperCAmelCase_					)

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

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

			lowerCamelCase     :			str		=    model.init_cache(decoder_input_ids.shape[0]		, UpperCAmelCase_		, UpperCAmelCase_					)
			lowerCamelCase     :			Dict		=    jnp.ones((decoder_input_ids.shape[0], max_decoder_length)		, dtype='i4'					)

			lowerCamelCase     :			Tuple		=    jnp.broadcast_to(
			    jnp.arange(decoder_input_ids.shape[-1] - 1					)[None, :]		, (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)		, )
			lowerCamelCase     :			int		=    model.decode(
			    decoder_input_ids[:, :-1]		, UpperCAmelCase_		, decoder_attention_mask=UpperCAmelCase_		, past_key_values=UpperCAmelCase_		, decoder_position_ids=UpperCAmelCase_		, )

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

			lowerCamelCase     :			Optional[Any]		=    model.decode(UpperCAmelCase_		, UpperCAmelCase_					)

			lowerCamelCase     :			Tuple		=    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 			_UpperCamelCase    ( self		, UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_					)		->							Union[str, Any]:
			lowerCamelCase     :			int		=    20
			lowerCamelCase     :			Tuple		=    model_class_name(UpperCAmelCase_					)

			lowerCamelCase     :			Optional[Any]		=    model.encode(inputs_dict['input_ids']					)

			lowerCamelCase	,    lowerCamelCase     :			Dict		=    (
			    inputs_dict['decoder_input_ids'],
			    inputs_dict['decoder_attention_mask'],
			)

			lowerCamelCase     :			Optional[Any]		=    jnp.concatenate(
			    [
			        decoder_attention_mask,
			        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])					),
			    ]		, axis=-1		, )

			lowerCamelCase     :			Optional[Any]		=    model.init_cache(decoder_input_ids.shape[0]		, UpperCAmelCase_		, UpperCAmelCase_					)
			lowerCamelCase     :			Tuple		=    jnp.broadcast_to(
			    jnp.arange(decoder_input_ids.shape[-1] - 1					)[None, :]		, (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)		, )

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

			lowerCamelCase     :			List[Any]		=    model.decode(UpperCAmelCase_		, UpperCAmelCase_		, decoder_attention_mask=UpperCAmelCase_					)

			lowerCamelCase     :			Tuple		=    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        _lowercase   (				unittest.TestCase		):
	lowercase_        =	9_9




	def 			_UpperCamelCase    ( self					)		->							Optional[Any]:
			lowerCamelCase     :			Optional[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		, )

			lowerCamelCase     :			List[Any]		=    input_ids.shape[0]
			lowerCamelCase     :			Optional[Any]		=    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 			_UpperCamelCase    ( self					)		->							Union[str, Any]:
			lowerCamelCase	,    lowerCamelCase	,    lowerCamelCase     :			List[Any]		=    self._get_config_and_data()
			lowerCamelCase     :			str		=    FlaxMBartForSequenceClassification(UpperCAmelCase_					)
			lowerCamelCase     :			Tuple		=    model(input_ids=UpperCAmelCase_		, decoder_input_ids=UpperCAmelCase_					)
			lowerCamelCase     :			Optional[int]		=    (batch_size, config.num_labels)
			self.assertEqual(outputs['logits'].shape		, UpperCAmelCase_					)




	def 			_UpperCamelCase    ( self					)		->							Any:
			lowerCamelCase	,    lowerCamelCase	,    lowerCamelCase     :			str		=    self._get_config_and_data()
			lowerCamelCase     :			Any		=    FlaxMBartForQuestionAnswering(UpperCAmelCase_					)
			lowerCamelCase     :			Any		=    model(input_ids=UpperCAmelCase_					)

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




	def 			_UpperCamelCase    ( self					)		->							Optional[Any]:
			lowerCamelCase	,    lowerCamelCase	,    lowerCamelCase     :			Optional[Any]		=    self._get_config_and_data()
			lowerCamelCase     :			List[str]		=    FlaxMBartForConditionalGeneration(UpperCAmelCase_					)
			lowerCamelCase     :			Optional[int]		=    lm_model(input_ids=UpperCAmelCase_					)
			lowerCamelCase     :			List[str]		=    (batch_size, input_ids.shape[1], config.vocab_size)
			self.assertEqual(outputs['logits'].shape		, UpperCAmelCase_					)




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




	def 			_UpperCamelCase    ( self					)		->							Tuple:
			lowerCamelCase     :			Optional[int]		=    np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]		, dtype=np.intaa					)
			lowerCamelCase     :			List[Any]		=    shift_tokens_right(UpperCAmelCase_		, 1					)
			lowerCamelCase     :			str		=    np.equal(UpperCAmelCase_		, 1					).astype(np.floataa					).sum()
			lowerCamelCase     :			Any		=    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        _lowercase   (				__UpperCAmelCase				,       unittest.TestCase				,       __UpperCAmelCase		):
	lowercase_        =	True
	lowercase_        =	(
	    (
	        FlaxMBartModel,
	        FlaxMBartForConditionalGeneration,
	        FlaxMBartForSequenceClassification,
	        FlaxMBartForQuestionAnswering,
	    )
	    if is_flax_available()
	    else ()
	)
	lowercase_        =	(FlaxMBartForConditionalGeneration,) if is_flax_available() else ()




	def 			_UpperCamelCase    ( self					)		->							Any:
			lowerCamelCase     :			Any		=    FlaxMBartModelTester(self					)




	def 			_UpperCamelCase    ( self					)		->							Optional[int]:
			lowerCamelCase	,    lowerCamelCase     :			int		=    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 			_UpperCamelCase    ( self					)		->							List[Any]:
			lowerCamelCase	,    lowerCamelCase     :			str		=    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 			_UpperCamelCase    ( self					)		->							List[Any]:
			lowerCamelCase	,    lowerCamelCase     :			Tuple		=    self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					with self.subTest(model_class.__name__					):
							lowerCamelCase     :			Union[str, Any]		=    self._prepare_for_class(UpperCAmelCase_		, UpperCAmelCase_					)
							lowerCamelCase     :			Dict		=    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     :			Dict		=    encode_jitted(**UpperCAmelCase_					).to_tuple()

							with self.subTest('JIT Disabled'					):
									with jax.disable_jit():
											lowerCamelCase     :			List[Any]		=    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 			_UpperCamelCase    ( self					)		->							str:
			lowerCamelCase	,    lowerCamelCase     :			Union[str, Any]		=    self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					with self.subTest(model_class.__name__					):
							lowerCamelCase     :			int		=    model_class(UpperCAmelCase_					)
							lowerCamelCase     :			Any		=    model.encode(inputs_dict['input_ids']		, inputs_dict['attention_mask']					)

							lowerCamelCase     :			Optional[int]		=    {
							    '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     :			int		=    decode_jitted(**UpperCAmelCase_					).to_tuple()

							with self.subTest('JIT Disabled'					):
									with jax.disable_jit():
											lowerCamelCase     :			Dict		=    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 			_UpperCamelCase    ( self					)		->							Union[str, Any]:
			for model_class_name in self.all_model_classes:
					lowerCamelCase     :			List[Any]		=    model_class_name.from_pretrained('facebook/mbart-large-cc25'		, from_pt=UpperCAmelCase_					)
					# FlaxMBartForSequenceClassification expects eos token in input_ids
					lowerCamelCase     :			str		=    np.ones((1, 1)					) * model.config.eos_token_id
					lowerCamelCase     :			int		=    model(UpperCAmelCase_					)
					self.assertIsNotNone(UpperCAmelCase_					)






@require_flax
@require_sentencepiece
@require_tokenizers
class        _lowercase   (				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 			_UpperCamelCase    ( self					)		->							Union[str, Any]:
			return AutoTokenizer.from_pretrained(self.model_name					)




	@cached_property
	def 			_UpperCamelCase    ( self					)		->							Optional[Any]:
			lowerCamelCase     :			Optional[int]		=    FlaxMBartForConditionalGeneration.from_pretrained(self.model_name		, from_pt=UpperCAmelCase_					)
			return model




	def 			_UpperCamelCase    ( self		, **UpperCAmelCase_					)		->							int:
			lowerCamelCase     :			List[str]		=    self.translate_src_text(**UpperCAmelCase_					)
			self.assertListEqual(self.expected_text		, UpperCAmelCase_					)




	def 			_UpperCamelCase    ( self		, **UpperCAmelCase_					)		->							Dict:
			lowerCamelCase     :			Optional[int]		=    self.tokenizer(self.src_text		, **UpperCAmelCase_		, return_tensors='np'					)
			lowerCamelCase     :			List[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=UpperCAmelCase_		, num_beams=2		, ).sequences
			lowerCamelCase     :			str		=    self.tokenizer.batch_decode(UpperCAmelCase_		, skip_special_tokens=UpperCAmelCase_					)
			return generated_words




	@slow
	def 			_UpperCamelCase    ( self					)		->							Union[str, Any]:
			self._assert_generated_batch_equal_expected()






