






"""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
       lowercase_								=	'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						,			):




  """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    snake_case						(				unittest.TestCase		):


   '''simple docstring'''

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

     '''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 							_SCREAMING_SNAKE_CASE							(      self					:	Dict    ):

     '''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(__snake_case,			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=__snake_case,			)
     __A		     =					prepare_mbart_inputs_dict(__snake_case,			__snake_case,			__snake_case    )
     return config, inputs_dict


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Any    ):

     '''simple docstring'''
     __A		     =					self.prepare_config_and_inputs()
     return config, inputs_dict


   def 							_SCREAMING_SNAKE_CASE							(      self					:	List[str],			_lowerCamelCase					:	str,			_lowerCamelCase					:	List[Any],			_lowerCamelCase					:	Optional[int]    ):

     '''simple docstring'''
     __A		     =					20
     __A		     =					model_class_name(__snake_case    )

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

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

     __A		     =					model.init_cache(decoder_input_ids.shape[0],			__snake_case,			__snake_case    )
     __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],			__snake_case,			decoder_attention_mask=__snake_case,			past_key_values=__snake_case,			decoder_position_ids=__snake_case,			)

     __A		     =					jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]],			dtype='''i4'''    )
     __A		     =					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,			)

     __A		     =					model.decode(__snake_case,			__snake_case    )

     __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 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[int],			_lowerCamelCase					:	int,			_lowerCamelCase					:	int,			_lowerCamelCase					:	int    ):

     '''simple docstring'''
     __A		     =					20
     __A		     =					model_class_name(__snake_case    )

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

     __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],			__snake_case,			__snake_case    )
     __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],			__snake_case,			decoder_attention_mask=__snake_case,			past_key_values=__snake_case,			decoder_position_ids=__snake_case,			)
     __A		     =					jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]],			dtype='''i4'''    )
     __A		     =					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,			)

     __A		     =					model.decode(__snake_case,			__snake_case,			decoder_attention_mask=__snake_case    )

     __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    snake_case						(				unittest.TestCase		):


   '''simple docstring'''

   A_		:      Tuple							=					99


   def 							_SCREAMING_SNAKE_CASE							(      self					:	List[str]    ):

     '''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 							_SCREAMING_SNAKE_CASE							(      self					:	Union[str, Any]    ):

     '''simple docstring'''
     __A		     =					self._get_config_and_data()
     __A		     =					FlaxMBartForSequenceClassification(__snake_case    )
     __A		     =					model(input_ids=__snake_case,			decoder_input_ids=__snake_case    )
     __A		     =					(batch_size, config.num_labels)
     self.assertEqual(outputs['''logits'''].shape,			__snake_case    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[Any]    ):

     '''simple docstring'''
     __A		     =					self._get_config_and_data()
     __A		     =					FlaxMBartForQuestionAnswering(__snake_case    )
     __A		     =					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'''
     __A		     =					self._get_config_and_data()
     __A		     =					FlaxMBartForConditionalGeneration(__snake_case    )
     __A		     =					lm_model(input_ids=__snake_case    )
     __A		     =					(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'''
     __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(__snake_case    )
     __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=__snake_case,			decoder_input_ids=__snake_case    )
     __A		     =					(*summary.shape, config.vocab_size)
     self.assertEqual(outputs['''logits'''].shape,			__snake_case    )






   def 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[int]    ):

     '''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(__snake_case,			1    )
     __A		     =					np.equal(__snake_case,			1    ).astype(np.floataa    ).sum()
     __A		     =					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    snake_case						(				__UpperCAmelCase  ,      unittest.TestCase  ,      __UpperCAmelCase		):


   '''simple docstring'''

   A_		:      str							=					True
   A_		:      Optional[Any]							=					(
       (
           FlaxMBartModel,
           FlaxMBartForConditionalGeneration,
           FlaxMBartForSequenceClassification,
           FlaxMBartForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )
   A_		:      int							=					(FlaxMBartForConditionalGeneration,) if is_flax_available() else ()


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Tuple    ):

     '''simple docstring'''
     __A		     =					FlaxMBartModelTester(self    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	int    ):

     '''simple docstring'''
     __A		     =					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'''
     __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(__snake_case,			__snake_case,			__snake_case    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	List[str]    ):

     '''simple docstring'''
     __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(__snake_case,			__snake_case    )
         __A		     =					model_class(__snake_case    )

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

         with self.subTest('''JIT Enabled'''    ):
           __A		     =					encode_jitted(**__snake_case    ).to_tuple()

         with self.subTest('''JIT Disabled'''    ):
           with jax.disable_jit():
             __A		     =					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'''
     __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(__snake_case    )
         __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					:	Dict,			_lowerCamelCase					:	Dict,			_lowerCamelCase					:	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'''    ):
           __A		     =					decode_jitted(**__snake_case    ).to_tuple()

         with self.subTest('''JIT Disabled'''    ):
           with jax.disable_jit():
             __A		     =					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:
       __A		     =					model_class_name.from_pretrained('''facebook/mbart-large-cc25''',			from_pt=__snake_case    )
       # FlaxMBartForSequenceClassification expects eos token in input_ids
       __A		     =					np.ones((1, 1)    ) * model.config.eos_token_id
       __A		     =					model(__snake_case    )
       self.assertIsNotNone(__snake_case    )





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


   '''simple docstring'''

   A_		:      str							=					[
       " UN Chief Says There Is No Military Solution in Syria",
   ]
   A_		:      List[Any]							=					[
       "Şeful ONU declară că nu există o soluţie militară în Siria",
   ]
   A_		:      Dict							=					"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'''
     __A		     =					FlaxMBartForConditionalGeneration.from_pretrained(self.model_name,			from_pt=__snake_case    )
     return model


   def 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[int],			**_lowerCamelCase					:	int    ):

     '''simple docstring'''
     __A		     =					self.translate_src_text(**__snake_case    )
     self.assertListEqual(self.expected_text,			__snake_case    )


   def 							_SCREAMING_SNAKE_CASE							(      self					:	str,			**_lowerCamelCase					:	Tuple    ):

     '''simple docstring'''
     __A		     =					self.tokenizer(self.src_text,			**__snake_case,			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=__snake_case,			num_beams=2,			).sequences
     __A		     =					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()











'''simple docstring'''






from __future__ import annotations

import copy
import os
import tempfile
from importlib import import_module
from math import isnan

from transformers import is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import _tf_gpu_memory_limit, require_tf, slow

from ..test_modeling_tf_common import ids_tensor


if is_tf_available():
		import numpy as np
		import tensorflow as tf

		from transformers import (
		    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
		    TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
		    TF_MODEL_FOR_MASKED_LM_MAPPING,
		    TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
		    TF_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING,
		    TF_MODEL_FOR_PRETRAINING_MAPPING,
		    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
		    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
		    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
		    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
		    TFSharedEmbeddings,
		)

		if _tf_gpu_memory_limit is not None:
				A :			Tuple							   =      tf.config.list_physical_devices('GPU')
				for gpu in gpus:
						# Restrict TensorFlow to only allocate x GB of memory on the GPUs
						try:
								tf.config.set_logical_device_configuration(
								    gpu, [tf.config.LogicalDeviceConfiguration(memory_limit=_tf_gpu_memory_limit)]
								)
								A :			Tuple							   =      tf.config.list_logical_devices('GPU')
								print('Logical GPUs', logical_gpus)
						except RuntimeError as e:
								# Virtual devices must be set before GPUs have been initialized
								print(e)







@require_tf
class       lowerCamelCase       :
					_SCREAMING_SNAKE_CASE				 =     None
					_SCREAMING_SNAKE_CASE				 =     ()
					_SCREAMING_SNAKE_CASE				 =     ()
					_SCREAMING_SNAKE_CASE				 =     True
					_SCREAMING_SNAKE_CASE				 =     True
					_SCREAMING_SNAKE_CASE				 =     True
					_SCREAMING_SNAKE_CASE				 =     False


					def  SCREAMING_SNAKE_CASE_		(       self						:	Union[str, Any]							,  __snake_case						:	Optional[Any]							,  __snake_case						:	int							,  __snake_case						:	Tuple=False							):





											'''simple docstring'''





											_snake_case: Tuple	       =		copy.deepcopy(__snake_case							)

											if model_class in get_values(__snake_case							):
																	_snake_case: Dict	       =		{
																	    k: tf.tile(tf.expand_dims(__snake_case							,  1							)							,  (1, self.model_tester.num_choices) + (1,) * (v.ndim - 1)							)
																	    if isinstance(__snake_case							,  tf.Tensor							) and v.ndim > 0
																	    else v
																	    for k, v in inputs_dict.items()
																	}

											if return_labels:
																	if model_class in get_values(__snake_case							):
																							_snake_case: Any	       =		tf.ones(self.model_tester.batch_size							,  dtype=tf.intaa							)
																	elif model_class in get_values(__snake_case							):
																							_snake_case: Dict	       =		tf.zeros(self.model_tester.batch_size							,  dtype=tf.intaa							)
																							_snake_case: Optional[Any]	       =		tf.zeros(self.model_tester.batch_size							,  dtype=tf.intaa							)
																	elif model_class in [
																	    *get_values(__snake_case							),
																	    *get_values(__snake_case							),
																	]:
																							_snake_case: int	       =		tf.zeros(self.model_tester.batch_size							,  dtype=tf.intaa							)
																	elif model_class in get_values(__snake_case							):
																							_snake_case: Optional[int]	       =		tf.zeros(self.model_tester.batch_size							,  dtype=tf.intaa							)
																	elif model_class in [
																	    *get_values(__snake_case							),
																	    *get_values(__snake_case							),
																	    *get_values(__snake_case							),
																	    *get_values(__snake_case							),
																	    *get_values(__snake_case							),
																	]:
																							_snake_case: Union[str, Any]	       =		tf.zeros(
																							    (self.model_tester.batch_size, self.model_tester.seq_length)							,  dtype=tf.intaa							)
											return inputs_dict


					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	Tuple							):





											'''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[:2]:
																	_snake_case: Tuple	       =		self._prepare_for_class(__snake_case							,  __snake_case							)
																	_snake_case: Dict	       =		model_class(__snake_case							)

																	@tf.function
																	def run_in_graph_mode():
																							return model(__snake_case							)

																	_snake_case: Dict	       =		run_in_graph_mode()
																	self.assertIsNotNone(__snake_case							)


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





											'''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[:2]:
																	_snake_case: Optional[Any]	       =		self._prepare_for_class(__snake_case							,  __snake_case							)
																	_snake_case: str	       =		model_class(__snake_case							)

																	@tf.function(experimental_compile=__snake_case							)
																	def run_in_graph_mode():
																							return model(__snake_case							)

																	_snake_case: str	       =		run_in_graph_mode()
																	self.assertIsNotNone(__snake_case							)


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





											'''simple docstring'''





											_snake_case       ,    _snake_case: List[Any]	       =		self.model_tester.prepare_config_and_inputs_for_common()
											for model_class in self.all_model_classes[:2]:
																	_snake_case: Dict	       =		model_class(__snake_case							)
																	if getattr(__snake_case							,  'hf_compute_loss'							,  __snake_case							):
																							# Test that model correctly compute the loss with kwargs
																							_snake_case: Tuple	       =		self._prepare_for_class(inputs_dict.copy()							,  __snake_case							,  return_labels=__snake_case							)
																							# Is there a better way to remove these decoder inputs?
																							_snake_case: Union[str, Any]	       =		{
																							    key: val
																							    for key, val in prepared_for_class.items()
																							    if key not in ('head_mask', 'decoder_head_mask', 'cross_attn_head_mask', 'decoder_input_ids')
																							}

																							_snake_case: List[str]	       =		{
																							    'labels',
																							    'label',
																							    'label_ids',
																							    'start_positions',
																							    'start_position',
																							    'end_positions',
																							    'end_position',
																							    'next_sentence_label',
																							}
																							_snake_case: Tuple	       =		possible_label_cols.intersection(set(__snake_case							)							)
																							self.assertGreater(len(__snake_case							)							,  0							,  msg='No matching label names found!'							)
																							_snake_case: Optional[Any]	       =		{key: val for key, val in prepared_for_class.items() if key in label_names}
																							_snake_case: Union[str, Any]	       =		{key: val for key, val in prepared_for_class.items() if key not in label_names}
																							self.assertGreater(len(__snake_case							)							,  0							)

																							# Make sure it works with XLA!
																							model.compile(optimizer=tf.keras.optimizers.SGD(0.0							)							,  jit_compile=__snake_case							)
																							# Make sure the model fits without crashing regardless of where we pass the labels
																							_snake_case: Optional[Any]	       =		model.fit(
																							    __snake_case							,  validation_data=__snake_case							,  steps_per_epoch=1							,  validation_steps=1							,  shuffle=__snake_case							,  verbose=0							,  )
																							_snake_case: str	       =		history.history['loss'][0]
																							self.assertTrue(not isnan(__snake_case							)							)
																							_snake_case: Tuple	       =		history.history['val_loss'][0]
																							self.assertTrue(not isnan(__snake_case							)							)

																							# Now test it with separate labels, to make sure that path works in XLA too.
																							_snake_case: Any	       =		model_class(__snake_case							)
																							model.compile(optimizer=tf.keras.optimizers.SGD(0.0							)							,  jit_compile=__snake_case							)
																							_snake_case: str	       =		model.fit(
																							    __snake_case							,  __snake_case							,  validation_data=(inputs_minus_labels, labels)							,  steps_per_epoch=1							,  validation_steps=1							,  shuffle=__snake_case							,  verbose=0							,  )

																							_snake_case: List[str]	       =		history.history['loss'][0]
																							self.assertTrue(not isnan(__snake_case							)							)
																							_snake_case: Optional[Any]	       =		history.history['val_loss'][0]
																							self.assertTrue(not isnan(__snake_case							)							)


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





											'''simple docstring'''





											_snake_case       ,    _snake_case: List[Any]	       =		self.model_tester.prepare_config_and_inputs_for_common()
											_snake_case: int	       =		True
											_snake_case: int	       =		True

											if hasattr(__snake_case							,  'use_cache'							):
																	_snake_case: Optional[Any]	       =		True

											_snake_case: Union[str, Any]	       =		getattr(self.model_tester							,  'encoder_seq_length'							,  self.model_tester.seq_length							)
											_snake_case: int	       =		getattr(self.model_tester							,  'key_length'							,  __snake_case							)

											for model_class in self.all_model_classes[:2]:
																	_snake_case: Tuple	       =		self._prepare_for_class(__snake_case							,  __snake_case							)
																	_snake_case: int	       =		model_class(__snake_case							)
																	model.build()
																	_snake_case: str	       =		len(model(__snake_case							)							)

																	for key in list(class_inputs_dict.keys()							):
																							# Remove keys not in the serving signature, as the SavedModel will not be compiled to deal with them
																							if key not in model.input_signature:
																													del class_inputs_dict[key]
																							# Check it's a tensor, in case the inputs dict has some bools in it too
																							elif isinstance(class_inputs_dict[key]							,  tf.Tensor							) and class_inputs_dict[key].dtype.is_integer:
																													_snake_case: Dict	       =		tf.cast(class_inputs_dict[key]							,  tf.intaa							)

																	if set(class_inputs_dict.keys()							) != set(model.input_signature.keys()							):
																							continue  # Some models have inputs that the preparation functions don't create, we skip those

																	with tempfile.TemporaryDirectory() as tmpdirname:
																							model.save_pretrained(__snake_case							,  saved_model=__snake_case							)
																							_snake_case: int	       =		os.path.join(__snake_case							,  'saved_model'							,  '1'							)
																							_snake_case: Optional[Any]	       =		tf.keras.models.load_model(__snake_case							)
																							_snake_case: Any	       =		model(__snake_case							)

																							if self.is_encoder_decoder:
																													_snake_case: str	       =		outputs['encoder_hidden_states']
																													_snake_case: Union[str, Any]	       =		outputs['encoder_attentions']
																							else:
																													_snake_case: Optional[int]	       =		outputs['hidden_states']
																													_snake_case: Dict	       =		outputs['attentions']

																							self.assertEqual(len(__snake_case							)							,  __snake_case							)

																							_snake_case: int	       =		getattr(
																							    self.model_tester							,  'expected_num_hidden_layers'							,  self.model_tester.num_hidden_layers + 1							)

																							self.assertEqual(len(__snake_case							)							,  __snake_case							)
																							self.assertListEqual(
																							    list(output_hidden_states[0].shape[-2:]							)							,  [self.model_tester.seq_length, self.model_tester.hidden_size]							,  )

																							self.assertEqual(len(__snake_case							)							,  self.model_tester.num_hidden_layers							)
																							self.assertListEqual(
																							    list(output_attentions[0].shape[-3:]							)							,  [self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length]							,  )


					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	str							):





											'''simple docstring'''





											tf.keras.mixed_precision.set_global_policy('mixed_float16'							)

											# try/finally block to ensure subsequent tests run in float32
											try:
																	_snake_case       ,    _snake_case: List[str]	       =		self.model_tester.prepare_config_and_inputs_for_common()
																	for model_class in self.all_model_classes[:2]:
																							_snake_case: Dict	       =		self._prepare_for_class(__snake_case							,  __snake_case							)
																							_snake_case: Any	       =		model_class(__snake_case							)
																							_snake_case: List[str]	       =		model(__snake_case							)

																							self.assertIsNotNone(__snake_case							)
											finally:
																	tf.keras.mixed_precision.set_global_policy('float32'							)


					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	Dict							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: Optional[int]	       =		self.model_tester.prepare_config_and_inputs_for_common()
											# head_mask and decoder_head_mask has different shapes than other input args
											if "head_mask" in inputs_dict:
																	del inputs_dict["head_mask"]
											if "decoder_head_mask" in inputs_dict:
																	del inputs_dict["decoder_head_mask"]
											if "cross_attn_head_mask" in inputs_dict:
																	del inputs_dict["cross_attn_head_mask"]
											_snake_case: List[str]	       =		{
											    module_member
											    for model_class in self.all_model_classes
											    for module in (import_module(model_class.__module__							),)
											    for module_member_name in dir(__snake_case							)
											    if module_member_name.endswith('MainLayer'							)
											    for module_member in (getattr(__snake_case							,  __snake_case							),)
											    if isinstance(__snake_case							,  __snake_case							)
											    and tf.keras.layers.Layer in module_member.__bases__
											    and getattr(__snake_case							,  '_keras_serializable'							,  __snake_case							)
											}

											for main_layer_class in tf_main_layer_classes:
																	# T5MainLayer needs an embed_tokens parameter when called without the inputs_embeds parameter
																	if "T5" in main_layer_class.__name__:
																							# Take the same values than in TFT5ModelTester for this shared layer
																							_snake_case: Optional[Any]	       =		TFSharedEmbeddings(self.model_tester.vocab_size							,  self.model_tester.hidden_size							,  name='shared'							)
																							_snake_case: str	       =		False
																							_snake_case: Tuple	       =		main_layer_class(__snake_case							,  embed_tokens=__snake_case							)
																	else:
																							_snake_case: List[Any]	       =		main_layer_class(__snake_case							)

																	_snake_case: Tuple	       =		{
																	    name: tf.keras.Input(tensor.shape[1:]							,  dtype=tensor.dtype							) for name, tensor in inputs_dict.items()
																	}

																	if hasattr(self.model_tester							,  'num_labels'							):
																							_snake_case: Dict	       =		self.model_tester.num_labels
																	else:
																							_snake_case: Tuple	       =		2

																	_snake_case: Tuple	       =		tf.data.Dataset.from_tensor_slices(
																	    (inputs_dict, np.ones((self.model_tester.batch_size, self.model_tester.seq_length, num_labels, 1)							))							).batch(1							)

																	_snake_case: List[str]	       =		main_layer(__snake_case							)[0]
																	_snake_case: List[Any]	       =		tf.keras.layers.Dense(__snake_case							,  activation='softmax'							,  name='outputs'							)(__snake_case							)
																	_snake_case: List[str]	       =		tf.keras.models.Model(inputs=__snake_case							,  outputs=[outputs]							)

																	model.compile(loss='binary_crossentropy'							,  optimizer='adam'							,  metrics=['binary_accuracy']							)
																	model.fit(__snake_case							,  epochs=1							)

																	with tempfile.TemporaryDirectory() as tmpdirname:
																							_snake_case: Tuple	       =		os.path.join(__snake_case							,  'keras_model.h5'							)
																							model.save(__snake_case							)
																							if "T5" in main_layer_class.__name__:
																													_snake_case: List[str]	       =		tf.keras.models.load_model(
																													    __snake_case							,  custom_objects={
																													        main_layer_class.__name__: main_layer_class,
																													        'TFSharedEmbeddings': TFSharedEmbeddings,
																													    }							,  )
																							else:
																													_snake_case: Optional[int]	       =		tf.keras.models.load_model(
																													    __snake_case							,  custom_objects={main_layer_class.__name__: main_layer_class}							)
																							assert isinstance(__snake_case							,  tf.keras.Model							)
																							model(__snake_case							)


					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	Tuple							):





											'''simple docstring'''





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

											for model_class in self.all_model_classes[:2]:
																	_snake_case: List[Any]	       =		model_class(__snake_case							)

																	_snake_case: Optional[int]	       =		copy.deepcopy(__snake_case							)

																	if not self.is_encoder_decoder:
																							_snake_case: Dict	       =		inputs['input_ids']
																							del inputs["input_ids"]
																	else:
																							_snake_case: Any	       =		inputs['input_ids']
																							_snake_case: List[str]	       =		inputs.get('decoder_input_ids'							,  __snake_case							)
																							del inputs["input_ids"]
																							inputs.pop('decoder_input_ids'							,  __snake_case							)

																	if not self.is_encoder_decoder:
																							_snake_case: str	       =		model.get_input_embeddings()(__snake_case							)
																	else:
																							_snake_case: Optional[int]	       =		model.get_input_embeddings()(__snake_case							)
																							_snake_case: List[str]	       =		model.get_input_embeddings()(__snake_case							)

																	_snake_case: str	       =		self._prepare_for_class(__snake_case							,  __snake_case							)

																	@tf.function
																	def run_in_graph_mode():
																							return model(__snake_case							)

																	_snake_case: Optional[Any]	       =		run_in_graph_mode()
																	self.assertIsNotNone(__snake_case							)


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





											'''simple docstring'''





											_snake_case: int	       =		[]
											if model.config.bos_token_id is not None:
																	special_tokens.append(model.config.bos_token_id							)
											if model.config.pad_token_id is not None:
																	special_tokens.append(model.config.pad_token_id							)
											if model.config.eos_token_id is not None:
																	special_tokens.append(model.config.eos_token_id							)

											# create random bad tokens that are not special tokens
											_snake_case: Tuple	       =		[]
											while len(__snake_case							) < num_bad_tokens:
																	_snake_case: Optional[Any]	       =		tf.squeeze(ids_tensor((1, 1)							,  self.model_tester.vocab_size							)							,  0							).numpy()[0]
																	if token not in special_tokens:
																							bad_tokens.append(__snake_case							)
											return bad_tokens


					def  SCREAMING_SNAKE_CASE_		(       self						:	Any							,  __snake_case						:	Tuple							):





											'''simple docstring'''





											for token_id in output_ids[0].numpy().tolist():
																	self.assertGreaterEqual(__snake_case							,  0							)
																	self.assertLess(__snake_case							,  self.model_tester.vocab_size							)




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





											'''simple docstring'''





											for bad_word_ids in bad_words_ids:
																	# for all slices in batch
																	for generated_ids_slice in generated_ids:
																							# for all word idx
																							for i in range(len(__snake_case							)							,  len(__snake_case							)							):
																													# if tokens match
																													if generated_ids_slice[i - len(__snake_case							) : i] == bad_word_ids:
																																			return True
											return False





