'''simple docstring'''

import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import OPTConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
      import os

      # The slow tests are often failing with OOM error on GPU
      # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
      # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
      __snake_case       :   int  		=    '''platform'''

      import jax
      import jax.numpy as jnp

      from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer



def lowerCamelCase__						(       A_       ,			A_       ,			A_=None       ,			A_=None       ):
 if attention_mask is None:
  UpperCAmelCase_				       =					np.where(input_ids != config.pad_token_id       ,			1       ,			0       )
 return {
     "input_ids": input_ids,
     "attention_mask": attention_mask,
 }
@require_flax
class 							lowercase_    :







    def __init__(	self       ,     UpperCamelCase__       ,     UpperCamelCase__=1_3       ,     UpperCamelCase__=7       ,     UpperCamelCase__=True       ,     UpperCamelCase__=False       ,     UpperCamelCase__=9_9       ,     UpperCamelCase__=1_6       ,     UpperCamelCase__=2       ,     UpperCamelCase__=4       ,     UpperCamelCase__=4       ,     UpperCamelCase__="gelu"       ,     UpperCamelCase__=0.1       ,     UpperCamelCase__=0.1       ,     UpperCamelCase__=2_0       ,     UpperCamelCase__=2       ,     UpperCamelCase__=1       ,     UpperCamelCase__=0       ,     UpperCamelCase__=1_6       ,     UpperCamelCase__=1_6       ,     UpperCamelCase__=0.02       ,     )				->       Dict:





     """simple docstring"""

     UpperCAmelCase_				       =					parent
     UpperCAmelCase_				       =					batch_size
     UpperCAmelCase_				       =					seq_length
     UpperCAmelCase_				       =					is_training
     UpperCAmelCase_				       =					use_labels
     UpperCAmelCase_				       =					vocab_size
     UpperCAmelCase_				       =					hidden_size
     UpperCAmelCase_				       =					num_hidden_layers
     UpperCAmelCase_				       =					num_attention_heads
     UpperCAmelCase_				       =					intermediate_size
     UpperCAmelCase_				       =					hidden_act
     UpperCAmelCase_				       =					hidden_dropout_prob
     UpperCAmelCase_				       =					attention_probs_dropout_prob
     UpperCAmelCase_				       =					max_position_embeddings
     UpperCAmelCase_				       =					eos_token_id
     UpperCAmelCase_				       =					pad_token_id
     UpperCAmelCase_				       =					bos_token_id
     UpperCAmelCase_				       =					embed_dim
     UpperCAmelCase_				       =					word_embed_proj_dim
     UpperCAmelCase_				       =					initializer_range
     UpperCAmelCase_				       =					False







    def   lowerCamelCase_    (	self				)				->       Dict:





     """simple docstring"""

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

     UpperCAmelCase_				       =					OPTConfig(
         vocab_size=self.vocab_size       ,     hidden_size=self.hidden_size       ,     num_hidden_layers=self.num_hidden_layers       ,     num_attention_heads=self.num_attention_heads       ,     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       ,     embed_dim=self.embed_dim       ,     is_encoder_decoder=__lowerCamelCase       ,     word_embed_proj_dim=self.word_embed_proj_dim       ,     initializer_range=self.initializer_range       ,     use_cache=__lowerCamelCase       ,     )
     UpperCAmelCase_				       =					prepare_opt_inputs_dict(__lowerCamelCase       ,     __lowerCamelCase				)
     return config, inputs_dict







    def   lowerCamelCase_    (	self				)				->       Dict:





     """simple docstring"""

     UpperCAmelCase_				       =					self.prepare_config_and_inputs()
     return config, inputs_dict







    def   lowerCamelCase_    (	self       ,     UpperCamelCase__       ,     UpperCamelCase__       ,     UpperCamelCase__				)				->       Dict:





     """simple docstring"""

     UpperCAmelCase_				       =					2_0
     UpperCAmelCase_				       =					model_class_name(__lowerCamelCase				)

     UpperCAmelCase_				       =					inputs_dict["input_ids"]
     UpperCAmelCase_				       =					inputs_dict["attention_mask"]

     UpperCAmelCase_				       =					model.init_cache(input_ids.shape[0]       ,     __lowerCamelCase				)
     UpperCAmelCase_				       =					jnp.ones((input_ids.shape[0], max_length)       ,     dtype="i4"				)

     UpperCAmelCase_				       =					jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1				)[None, :]       ,     (input_ids.shape[0], input_ids.shape[-1] - 1)       ,     )
     UpperCAmelCase_				       =					model(
         input_ids[:, :-1]       ,     attention_mask=__lowerCamelCase       ,     past_key_values=__lowerCamelCase       ,     position_ids=__lowerCamelCase       ,     )

     UpperCAmelCase_				       =					jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,     dtype="i4"				)
     UpperCAmelCase_				       =					model(
         input_ids[:, -1:]       ,     attention_mask=__lowerCamelCase       ,     past_key_values=outputs_cache.past_key_values       ,     position_ids=__lowerCamelCase       ,     )

     UpperCAmelCase_				       =					model(__lowerCamelCase				)

     UpperCAmelCase_				       =					np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])				)				)
     self.parent.assertTrue(diff < 1e-3       ,     msg=F"""Max diff is {diff}"""				)







    def   lowerCamelCase_    (	self       ,     UpperCamelCase__       ,     UpperCamelCase__       ,     UpperCamelCase__				)				->       Union[str, Any]:





     """simple docstring"""

     UpperCAmelCase_				       =					2_0
     UpperCAmelCase_				       =					model_class_name(__lowerCamelCase				)

     UpperCAmelCase_				       =					(
         inputs_dict["input_ids"],
         inputs_dict["attention_mask"],
     )

     UpperCAmelCase_				       =					jnp.concatenate(
         [
             attention_mask,
             jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])				),
         ]       ,     axis=-1       ,     )

     UpperCAmelCase_				       =					model.init_cache(input_ids.shape[0]       ,     __lowerCamelCase				)
     UpperCAmelCase_				       =					jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1				)[None, :]       ,     (input_ids.shape[0], input_ids.shape[-1] - 1)       ,     )

     UpperCAmelCase_				       =					model(
         input_ids[:, :-1]       ,     attention_mask=__lowerCamelCase       ,     past_key_values=__lowerCamelCase       ,     position_ids=__lowerCamelCase       ,     )
     UpperCAmelCase_				       =					jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,     dtype="i4"				)
     UpperCAmelCase_				       =					model(
         input_ids[:, -1:]       ,     past_key_values=outputs_cache.past_key_values       ,     attention_mask=__lowerCamelCase       ,     position_ids=__lowerCamelCase       ,     )

     UpperCAmelCase_				       =					model(__lowerCamelCase       ,     attention_mask=__lowerCamelCase				)

     UpperCAmelCase_				       =					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_    (     _A ,  unittest.TestCase ,  _A						):
    a_             =   (FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
    a_             =   () if is_flax_available() else ()







    def   lowerCamelCase_    (	self				)				->       str:





     """simple docstring"""

     UpperCAmelCase_				       =					FlaxOPTModelTester(self				)







    def   lowerCamelCase_    (	self				)				->       Optional[Any]:





     """simple docstring"""

     UpperCAmelCase_				       =					self.model_tester.prepare_config_and_inputs()
     for model_class in self.all_model_classes:
      self.model_tester.check_use_cache_forward(__lowerCamelCase       ,     __lowerCamelCase       ,     __lowerCamelCase				)







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





     """simple docstring"""

     UpperCAmelCase_				       =					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(__lowerCamelCase       ,     __lowerCamelCase       ,     __lowerCamelCase				)







    @slow
    def   lowerCamelCase_    (	self				)				->       Optional[Any]:





     """simple docstring"""

     for model_class_name in self.all_model_classes:
      UpperCAmelCase_				       =					model_class_name.from_pretrained("facebook/opt-125m"				)
      UpperCAmelCase_				       =					np.ones((1, 1)				) * model.config.eos_token_id
      UpperCAmelCase_				       =					model(__lowerCamelCase				)
      self.assertIsNotNone(__lowerCamelCase				)


@require_sentencepiece
@require_flax
class 							lowercase_    (     unittest.TestCase						):







    @slow
    def   lowerCamelCase_    (	self				)				->       List[str]:





     """simple docstring"""

     UpperCAmelCase_				       =					FlaxOPTModel.from_pretrained("facebook/opt-350m"				)
     UpperCAmelCase_				       =					jnp.array([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]]				)
     UpperCAmelCase_				       =					model(input_ids=__lowerCamelCase				).last_hidden_state
     UpperCAmelCase_				       =					(1, 1_1, 5_1_2)
     self.assertEqual(output.shape       ,     __lowerCamelCase				)
     UpperCAmelCase_				       =					jnp.array(
         [[-0.2867, -1.9256, -0.3062], [-1.2711, -0.1337, -0.1897], [0.4109, 0.1187, -1.3142]]				)
     self.assertTrue(jnp.allclose(output[:, :3, :3]       ,     __lowerCamelCase       ,     atol=4e-2				)				)


@require_flax
@slow
class 							lowercase_    (     unittest.TestCase						):







    def   lowerCamelCase_    (	self				)				->       Any:





     """simple docstring"""

     super().setUp()
     UpperCAmelCase_				       =					"facebook/opt-350m"







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





     """simple docstring"""

     UpperCAmelCase_				       =					FlaxOPTForCausalLM.from_pretrained(self.path_model				)
     UpperCAmelCase_				       =					GPTaTokenizer.from_pretrained(self.path_model				)

     UpperCAmelCase_				       =					[
         "Today is a beautiful day and I want to",
         "In the city of",
         "Paris is the capital of France and",
         "Computers and mobile phones have taken",
     ]
     # verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
     UpperCAmelCase_				       =					tokenizer(__lowerCamelCase       ,     return_tensors="jax"       ,     padding=__lowerCamelCase       ,     add_special_tokens=__lowerCamelCase				)
     UpperCAmelCase_				       =					model(inputs.input_ids       ,     attention_mask=inputs.attention_mask				)[0].mean(axis=-1				)
     UpperCAmelCase_				       =					jnp.array(
         [
             [1.3851, -1_3.8_9_2_3, -1_0.5_2_2_9, -1_0.7_5_3_3, -0.2309, -1_0.2_3_8_4, -0.5365, -9.0947, -5.1670],
             [-4.7073, -1_0.6_2_7_6, -3.9415, -2_1.5_2_4_2, -0.2822, -0.2822, -0.2822, -0.2822, -0.2822],
             [0.6247, -3.4229, -8.9179, -1.4297, -1_4.1_6_5_0, 1.4146, -9.0218, -0.2703, -0.2703],
             [6.4783, -1.9913, -1_0.7_9_2_6, -2.3336, 1.5092, -0.9974, -6.8213, 1.3477, 1.3477],
         ]				)
     self.assertTrue(jnp.allclose(__lowerCamelCase       ,     __lowerCamelCase       ,     atol=4e-2				)				)

     UpperCAmelCase_				       =					jax.jit(__lowerCamelCase				)
     UpperCAmelCase_				       =					model(inputs.input_ids       ,     attention_mask=inputs.attention_mask				)[0].mean(axis=-1				)
     self.assertTrue(jnp.allclose(__lowerCamelCase       ,     __lowerCamelCase       ,     atol=4e-2				)				)


@require_flax
@slow
class 							lowercase_    (     unittest.TestCase						):







    @property
    def   lowerCamelCase_    (	self				)				->       str:





     """simple docstring"""

     return [
         "Today is a beautiful day and I want",
         "In the city of",
         "Paris is the capital of France and",
         "Computers and mobile phones have taken",
     ]







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





     """simple docstring"""

     UpperCAmelCase_				       =					"facebook/opt-125m"

     UpperCAmelCase_				       =					[
         "Today is a beautiful day and I want to",
         "In the city of New York, the city",
         "Paris is the capital of France and the capital",
         "Computers and mobile phones have taken over the",
     ]

     UpperCAmelCase_				       =					[]

     UpperCAmelCase_				       =					FlaxOPTForCausalLM.from_pretrained(__lowerCamelCase				)
     UpperCAmelCase_				       =					GPTaTokenizer.from_pretrained(__lowerCamelCase				)

     for prompt in self.prompts:
      UpperCAmelCase_				       =					tokenizer(__lowerCamelCase       ,     return_tensors="jax"				).input_ids

      UpperCAmelCase_				       =					model.generate(__lowerCamelCase       ,     max_length=1_0				)
      UpperCAmelCase_				       =					generated_ids[0]

      UpperCAmelCase_				       =					tokenizer.batch_decode(__lowerCamelCase       ,     skip_special_tokens=__lowerCamelCase				)
      predicted_outputs += generated_string

     self.assertListEqual(__lowerCamelCase       ,     __lowerCamelCase				)







    def   lowerCamelCase_    (	self				)				->       List[str]:





     """simple docstring"""

     UpperCAmelCase_				       =					"facebook/opt-350m"

     UpperCAmelCase_				       =					[
         "Today is a beautiful day and I want to",
         "In the city of San Francisco, the city",
         "Paris is the capital of France and the capital",
         "Computers and mobile phones have taken over the",
     ]

     UpperCAmelCase_				       =					[]
     UpperCAmelCase_				       =					FlaxOPTForCausalLM.from_pretrained(__lowerCamelCase				)
     UpperCAmelCase_				       =					GPTaTokenizer.from_pretrained(__lowerCamelCase				)

     for prompt in self.prompts:
      UpperCAmelCase_				       =					tokenizer(__lowerCamelCase       ,     return_tensors="jax"				).input_ids

      UpperCAmelCase_				       =					model.generate(__lowerCamelCase       ,     max_length=1_0				)
      UpperCAmelCase_				       =					generated_ids[0]

      UpperCAmelCase_				       =					tokenizer.batch_decode(__lowerCamelCase       ,     skip_special_tokens=__lowerCamelCase				)
      predicted_outputs += generated_string

     self.assertListEqual(__lowerCamelCase       ,     __lowerCamelCase				)







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





     """simple docstring"""

     UpperCAmelCase_				       =					"facebook/opt-125m"
     UpperCAmelCase_				       =					[
         "Today is a beautiful day and I want to thank",
         "In the city of Rome Canaver Canaver Canaver Canaver",
     ]
     UpperCAmelCase_				       =					FlaxOPTForCausalLM.from_pretrained(__lowerCamelCase				)
     UpperCAmelCase_				       =					GPTaTokenizer.from_pretrained(__lowerCamelCase				)
     UpperCAmelCase_				       =					tokenizer(
         [
             "Today is a beautiful day and I want to",
             "In the city of",
         ]       ,     return_tensors="jax"       ,     padding=__lowerCamelCase       ,     )

     UpperCAmelCase_				       =					jax.jit(model.generate				)

     UpperCAmelCase_				       =					jit_generate(inputs["input_ids"]       ,     attention_mask=inputs["attention_mask"]				).sequences

     UpperCAmelCase_				       =					tokenizer.batch_decode(__lowerCamelCase       ,     skip_special_tokens=__lowerCamelCase				)

     self.assertIsNotNone(__lowerCamelCase       ,     __lowerCamelCase				)







    def   lowerCamelCase_    (	self				)				->       Dict:





     """simple docstring"""

     UpperCAmelCase_				       =					"facebook/opt-350m"

     UpperCAmelCase_				       =					GPTaTokenizer.from_pretrained(__lowerCamelCase				)
     UpperCAmelCase_				       =					FlaxOPTForCausalLM.from_pretrained(__lowerCamelCase				)

     UpperCAmelCase_				       =					"left"

     # use different length sentences to test batching
     UpperCAmelCase_				       =					[
         "Hello, my dog is a little",
         "Today, I",
     ]

     UpperCAmelCase_				       =					tokenizer(__lowerCamelCase       ,     return_tensors="jax"       ,     padding=__lowerCamelCase				)
     UpperCAmelCase_				       =					inputs["input_ids"]

     UpperCAmelCase_				       =					model.generate(input_ids=__lowerCamelCase       ,     attention_mask=inputs["attention_mask"]       ,     trace=__lowerCamelCase				)

     UpperCAmelCase_				       =					tokenizer(sentences[0]       ,     return_tensors="jax"				).input_ids
     UpperCAmelCase_				       =					model.generate(input_ids=__lowerCamelCase				)

     UpperCAmelCase_				       =					inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].sum()
     UpperCAmelCase_				       =					tokenizer(sentences[1]       ,     return_tensors="jax"				).input_ids
     UpperCAmelCase_				       =					model.generate(input_ids=__lowerCamelCase       ,     max_length=model.config.max_length - num_paddings				)

     UpperCAmelCase_				       =					tokenizer.batch_decode(outputs[0]       ,     skip_special_tokens=__lowerCamelCase				)
     UpperCAmelCase_				       =					tokenizer.decode(output_non_padded[0][0]       ,     skip_special_tokens=__lowerCamelCase				)
     UpperCAmelCase_				       =					tokenizer.decode(output_padded[0][0]       ,     skip_special_tokens=__lowerCamelCase				)

     UpperCAmelCase_				       =					[
         "Hello, my dog is a little bit of a dork.\nI'm a little bit",
         "Today, I was in the middle of a conversation with a friend about the",
     ]
     self.assertListEqual(__lowerCamelCase       ,     __lowerCamelCase				)
     self.assertListEqual(__lowerCamelCase       ,     [non_padded_sentence, padded_sentence]				)

'''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:
            __snake_case       :   Any  		=    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)]
                        )
                        __snake_case       :   Optional[Any]  		=    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 							lowercase_    :
    a_             =   None
    a_             =   ()
    a_             =   ()
    a_             =   True
    a_             =   True
    a_             =   True
    a_             =   False







    def   lowerCamelCase_    (	self       ,     UpperCamelCase__       ,     UpperCamelCase__       ,     UpperCamelCase__=False				)				->       dict:





     """simple docstring"""

     UpperCAmelCase_				       =					copy.deepcopy(UpperCamelCase__				)

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

     if return_labels:
      if model_class in get_values(UpperCamelCase__				):
       UpperCAmelCase_				       =					tf.ones(self.model_tester.batch_size       ,     dtype=tf.intaa				)
      elif model_class in get_values(UpperCamelCase__				):
       UpperCAmelCase_				       =					tf.zeros(self.model_tester.batch_size       ,     dtype=tf.intaa				)
       UpperCAmelCase_				       =					tf.zeros(self.model_tester.batch_size       ,     dtype=tf.intaa				)
      elif model_class in [
          *get_values(UpperCamelCase__				),
          *get_values(UpperCamelCase__				),
      ]:
       UpperCAmelCase_				       =					tf.zeros(self.model_tester.batch_size       ,     dtype=tf.intaa				)
      elif model_class in get_values(UpperCamelCase__				):
       UpperCAmelCase_				       =					tf.zeros(self.model_tester.batch_size       ,     dtype=tf.intaa				)
      elif model_class in [
          *get_values(UpperCamelCase__				),
          *get_values(UpperCamelCase__				),
          *get_values(UpperCamelCase__				),
          *get_values(UpperCamelCase__				),
          *get_values(UpperCamelCase__				),
      ]:
       UpperCAmelCase_				       =					tf.zeros(
           (self.model_tester.batch_size, self.model_tester.seq_length)       ,     dtype=tf.intaa				)
     return inputs_dict







    @slow
    def   lowerCamelCase_    (	self				)				->       Optional[Any]:





     """simple docstring"""

     UpperCAmelCase_       ,       UpperCAmelCase_				       =					self.model_tester.prepare_config_and_inputs_for_common()
     for model_class in self.all_model_classes[:2]:
      UpperCAmelCase_				       =					self._prepare_for_class(UpperCamelCase__       ,     UpperCamelCase__				)
      UpperCAmelCase_				       =					model_class(UpperCamelCase__				)

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

      UpperCAmelCase_				       =					run_in_graph_mode()
      self.assertIsNotNone(UpperCamelCase__				)







    @slow
    def   lowerCamelCase_    (	self				)				->       Optional[int]:





     """simple docstring"""

     UpperCAmelCase_       ,       UpperCAmelCase_				       =					self.model_tester.prepare_config_and_inputs_for_common()
     for model_class in self.all_model_classes[:2]:
      UpperCAmelCase_				       =					self._prepare_for_class(UpperCamelCase__       ,     UpperCamelCase__				)
      UpperCAmelCase_				       =					model_class(UpperCamelCase__				)

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

      UpperCAmelCase_				       =					run_in_graph_mode()
      self.assertIsNotNone(UpperCamelCase__				)







    @slow
    def   lowerCamelCase_    (	self				)				->       int:





     """simple docstring"""

     UpperCAmelCase_       ,       UpperCAmelCase_				       =					self.model_tester.prepare_config_and_inputs_for_common()
     for model_class in self.all_model_classes[:2]:
      UpperCAmelCase_				       =					model_class(UpperCamelCase__				)
      if getattr(UpperCamelCase__       ,     "hf_compute_loss"       ,     UpperCamelCase__				):
       # Test that model correctly compute the loss with kwargs
       UpperCAmelCase_				       =					self._prepare_for_class(inputs_dict.copy()       ,     UpperCamelCase__       ,     return_labels=UpperCamelCase__				)
       # Is there a better way to remove these decoder inputs?
       UpperCAmelCase_				       =					{
           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")
       }

       UpperCAmelCase_				       =					{
           "labels",
           "label",
           "label_ids",
           "start_positions",
           "start_position",
           "end_positions",
           "end_position",
           "next_sentence_label",
       }
       UpperCAmelCase_				       =					possible_label_cols.intersection(set(UpperCamelCase__				)				)
       self.assertGreater(len(UpperCamelCase__				)       ,     0       ,     msg="No matching label names found!"				)
       UpperCAmelCase_				       =					{key: val for key, val in prepared_for_class.items() if key in label_names}
       UpperCAmelCase_				       =					{key: val for key, val in prepared_for_class.items() if key not in label_names}
       self.assertGreater(len(UpperCamelCase__				)       ,     0				)

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

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

       UpperCAmelCase_				       =					history.history["loss"][0]
       self.assertTrue(not isnan(UpperCamelCase__				)				)
       UpperCAmelCase_				       =					history.history["val_loss"][0]
       self.assertTrue(not isnan(UpperCamelCase__				)				)







    @slow
    def   lowerCamelCase_    (	self				)				->       Tuple:





     """simple docstring"""

     UpperCAmelCase_       ,       UpperCAmelCase_				       =					self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase_				       =					True
     UpperCAmelCase_				       =					True

     if hasattr(UpperCamelCase__       ,     "use_cache"				):
      UpperCAmelCase_				       =					True

     UpperCAmelCase_				       =					getattr(self.model_tester       ,     "encoder_seq_length"       ,     self.model_tester.seq_length				)
     UpperCAmelCase_				       =					getattr(self.model_tester       ,     "key_length"       ,     UpperCamelCase__				)

     for model_class in self.all_model_classes[:2]:
      UpperCAmelCase_				       =					self._prepare_for_class(UpperCamelCase__       ,     UpperCamelCase__				)
      UpperCAmelCase_				       =					model_class(UpperCamelCase__				)
      model.build()
      UpperCAmelCase_				       =					len(model(UpperCamelCase__				)				)

      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:
        UpperCAmelCase_				       =					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(UpperCamelCase__       ,     saved_model=UpperCamelCase__				)
       UpperCAmelCase_				       =					os.path.join(UpperCamelCase__       ,     "saved_model"       ,     "1"				)
       UpperCAmelCase_				       =					tf.keras.models.load_model(UpperCamelCase__				)
       UpperCAmelCase_				       =					model(UpperCamelCase__				)

       if self.is_encoder_decoder:
        UpperCAmelCase_				       =					outputs["encoder_hidden_states"]
        UpperCAmelCase_				       =					outputs["encoder_attentions"]
       else:
        UpperCAmelCase_				       =					outputs["hidden_states"]
        UpperCAmelCase_				       =					outputs["attentions"]

       self.assertEqual(len(UpperCamelCase__				)       ,     UpperCamelCase__				)

       UpperCAmelCase_				       =					getattr(
           self.model_tester       ,     "expected_num_hidden_layers"       ,     self.model_tester.num_hidden_layers + 1				)

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

       self.assertEqual(len(UpperCamelCase__				)       ,     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   lowerCamelCase_    (	self				)				->       Optional[int]:





     """simple docstring"""

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

     # try/finally block to ensure subsequent tests run in float32
     try:
      UpperCAmelCase_       ,       UpperCAmelCase_				       =					self.model_tester.prepare_config_and_inputs_for_common()
      for model_class in self.all_model_classes[:2]:
       UpperCAmelCase_				       =					self._prepare_for_class(UpperCamelCase__       ,     UpperCamelCase__				)
       UpperCAmelCase_				       =					model_class(UpperCamelCase__				)
       UpperCAmelCase_				       =					model(UpperCamelCase__				)

       self.assertIsNotNone(UpperCamelCase__				)
     finally:
      tf.keras.mixed_precision.set_global_policy("float32"				)







    @slow
    def   lowerCamelCase_    (	self				)				->       Dict:





     """simple docstring"""

     UpperCAmelCase_       ,       UpperCAmelCase_				       =					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"]
     UpperCAmelCase_				       =					{
         module_member
         for model_class in self.all_model_classes
         for module in (import_module(model_class.__module__				),)
         for module_member_name in dir(UpperCamelCase__				)
         if module_member_name.endswith("MainLayer"				)
         for module_member in (getattr(UpperCamelCase__       ,     UpperCamelCase__				),)
         if isinstance(UpperCamelCase__       ,     UpperCamelCase__				)
         and tf.keras.layers.Layer in module_member.__bases__
         and getattr(UpperCamelCase__       ,     "_keras_serializable"       ,     UpperCamelCase__				)
     }

     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
       UpperCAmelCase_				       =					TFSharedEmbeddings(self.model_tester.vocab_size       ,     self.model_tester.hidden_size       ,     name="shared"				)
       UpperCAmelCase_				       =					False
       UpperCAmelCase_				       =					main_layer_class(UpperCamelCase__       ,     embed_tokens=UpperCamelCase__				)
      else:
       UpperCAmelCase_				       =					main_layer_class(UpperCamelCase__				)

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

      if hasattr(self.model_tester       ,     "num_labels"				):
       UpperCAmelCase_				       =					self.model_tester.num_labels
      else:
       UpperCAmelCase_				       =					2

      UpperCAmelCase_				       =					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				)

      UpperCAmelCase_				       =					main_layer(UpperCamelCase__				)[0]
      UpperCAmelCase_				       =					tf.keras.layers.Dense(UpperCamelCase__       ,     activation="softmax"       ,     name="outputs"				)(UpperCamelCase__				)
      UpperCAmelCase_				       =					tf.keras.models.Model(inputs=UpperCamelCase__       ,     outputs=[outputs]				)

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

      with tempfile.TemporaryDirectory() as tmpdirname:
       UpperCAmelCase_				       =					os.path.join(UpperCamelCase__       ,     "keras_model.h5"				)
       model.save(UpperCamelCase__				)
       if "T5" in main_layer_class.__name__:
        UpperCAmelCase_				       =					tf.keras.models.load_model(
            UpperCamelCase__       ,     custom_objects={
                main_layer_class.__name__: main_layer_class,
                "TFSharedEmbeddings": TFSharedEmbeddings,
            }       ,     )
       else:
        UpperCAmelCase_				       =					tf.keras.models.load_model(
            UpperCamelCase__       ,     custom_objects={main_layer_class.__name__: main_layer_class}				)
       assert isinstance(UpperCamelCase__       ,     tf.keras.Model				)
       model(UpperCamelCase__				)







    @slow
    def   lowerCamelCase_    (	self				)				->       str:





     """simple docstring"""

     UpperCAmelCase_       ,       UpperCAmelCase_				       =					self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes[:2]:
      UpperCAmelCase_				       =					model_class(UpperCamelCase__				)

      UpperCAmelCase_				       =					copy.deepcopy(UpperCamelCase__				)

      if not self.is_encoder_decoder:
       UpperCAmelCase_				       =					inputs["input_ids"]
       del inputs["input_ids"]
      else:
       UpperCAmelCase_				       =					inputs["input_ids"]
       UpperCAmelCase_				       =					inputs.get("decoder_input_ids"       ,     UpperCamelCase__				)
       del inputs["input_ids"]
       inputs.pop("decoder_input_ids"       ,     UpperCamelCase__				)

      if not self.is_encoder_decoder:
       UpperCAmelCase_				       =					model.get_input_embeddings()(UpperCamelCase__				)
      else:
       UpperCAmelCase_				       =					model.get_input_embeddings()(UpperCamelCase__				)
       UpperCAmelCase_				       =					model.get_input_embeddings()(UpperCamelCase__				)

      UpperCAmelCase_				       =					self._prepare_for_class(UpperCamelCase__       ,     UpperCamelCase__				)

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

      UpperCAmelCase_				       =					run_in_graph_mode()
      self.assertIsNotNone(UpperCamelCase__				)







    def   lowerCamelCase_    (	self       ,     UpperCamelCase__       ,     UpperCamelCase__				)				->       Dict:





     """simple docstring"""

     UpperCAmelCase_				       =					[]
     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
     UpperCAmelCase_				       =					[]
     while len(UpperCamelCase__				) < num_bad_tokens:
      UpperCAmelCase_				       =					tf.squeeze(ids_tensor((1, 1)       ,     self.model_tester.vocab_size				)       ,     0				).numpy()[0]
      if token not in special_tokens:
       bad_tokens.append(UpperCamelCase__				)
     return bad_tokens







    def   lowerCamelCase_    (	self       ,     UpperCamelCase__				)				->       List[str]:





     """simple docstring"""

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







    def   lowerCamelCase_    (	self       ,     UpperCamelCase__       ,     UpperCamelCase__				)				->       Any:





     """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(UpperCamelCase__				)       ,     len(UpperCamelCase__				)				):
        # if tokens match
        if generated_ids_slice[i - len(UpperCamelCase__				) : i] == bad_word_ids:
         return True
     return False
