


'''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:
  lowerCamelCase							:       Union[str, 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)]
    )
    lowerCamelCase							:       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   __lowerCAmelCase						:





      '''simple docstring'''






      lowerCAmelCase__ :					Optional[int]   =							None
      lowerCAmelCase__ :					Optional[Any]   =							()
      lowerCAmelCase__ :					int   =							()
      lowerCAmelCase__ :					List[Any]   =							True
      lowerCAmelCase__ :					str   =							True
      lowerCAmelCase__ :					int   =							True
      lowerCAmelCase__ :					Optional[int]   =							False



      def 				UpperCamelCase__		(self   :   Tuple       ,						UpperCamelCase   :   Optional[Any]       ,						UpperCamelCase   :   List[str]       ,						UpperCamelCase   :   List[Any]=False			):

             '''simple docstring'''







             lowercase__				      =			copy.deepcopy(UpperCamelCase			)

             if model_class in get_values(UpperCamelCase			):
                    lowercase__				      =			{
                        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			):
                           lowercase__				      =			tf.ones(self.model_tester.batch_size       ,						dtype=tf.intaa			)
                    elif model_class in get_values(UpperCamelCase			):
                           lowercase__				      =			tf.zeros(self.model_tester.batch_size       ,						dtype=tf.intaa			)
                           lowercase__				      =			tf.zeros(self.model_tester.batch_size       ,						dtype=tf.intaa			)
                    elif model_class in [
                        *get_values(UpperCamelCase			),
                        *get_values(UpperCamelCase			),
                    ]:
                           lowercase__				      =			tf.zeros(self.model_tester.batch_size       ,						dtype=tf.intaa			)
                    elif model_class in get_values(UpperCamelCase			):
                           lowercase__				      =			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			),
                    ]:
                           lowercase__				      =			tf.zeros(
                               (self.model_tester.batch_size, self.model_tester.seq_length)       ,						dtype=tf.intaa			)
             return inputs_dict



      @slow
      def 				UpperCamelCase__		(self   :   int			):

             '''simple docstring'''







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

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

                    lowercase__				      =			run_in_graph_mode()
                    self.assertIsNotNone(UpperCamelCase			)



      @slow
      def 				UpperCamelCase__		(self   :   Tuple			):

             '''simple docstring'''







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

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

                    lowercase__				      =			run_in_graph_mode()
                    self.assertIsNotNone(UpperCamelCase			)



      @slow
      def 				UpperCamelCase__		(self   :   List[str]			):

             '''simple docstring'''







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

                           lowercase__				      =			{
                               '''labels''',
                               '''label''',
                               '''label_ids''',
                               '''start_positions''',
                               '''start_position''',
                               '''end_positions''',
                               '''end_position''',
                               '''next_sentence_label''',
                           }
                           lowercase__				      =			possible_label_cols.intersection(set(UpperCamelCase			)			)
                           self.assertGreater(len(UpperCamelCase			)       ,						0       ,						msg='''No matching label names found!'''			)
                           lowercase__				      =			{key: val for key, val in prepared_for_class.items() if key in label_names}
                           lowercase__				      =			{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
                           lowercase__				      =			model.fit(
                               UpperCamelCase       ,						validation_data=UpperCamelCase       ,						steps_per_epoch=1       ,						validation_steps=1       ,						shuffle=UpperCamelCase       ,						verbose=0       ,						)
                           lowercase__				      =			history.history['''loss'''][0]
                           self.assertTrue(not isnan(UpperCamelCase			)			)
                           lowercase__				      =			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.
                           lowercase__				      =			model_class(UpperCamelCase			)
                           model.compile(optimizer=tf.keras.optimizers.SGD(0.0			)       ,						jit_compile=UpperCamelCase			)
                           lowercase__				      =			model.fit(
                               UpperCamelCase       ,						UpperCamelCase       ,						validation_data=(inputs_minus_labels, labels)       ,						steps_per_epoch=1       ,						validation_steps=1       ,						shuffle=UpperCamelCase       ,						verbose=0       ,						)

                           lowercase__				      =			history.history['''loss'''][0]
                           self.assertTrue(not isnan(UpperCamelCase			)			)
                           lowercase__				      =			history.history['''val_loss'''][0]
                           self.assertTrue(not isnan(UpperCamelCase			)			)



      @slow
      def 				UpperCamelCase__		(self   :   str			):

             '''simple docstring'''







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

             if hasattr(UpperCamelCase       ,						'''use_cache'''			):
                    lowercase__				      =			True

             lowercase__				      =			getattr(self.model_tester       ,						'''encoder_seq_length'''       ,						self.model_tester.seq_length			)
             lowercase__				      =			getattr(self.model_tester       ,						'''key_length'''       ,						UpperCamelCase			)

             for model_class in self.all_model_classes[:2]:
                    lowercase__				      =			self._prepare_for_class(UpperCamelCase       ,						UpperCamelCase			)
                    lowercase__				      =			model_class(UpperCamelCase			)
                    model.build()
                    lowercase__				      =			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:
                                  lowercase__				      =			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			)
                           lowercase__				      =			os.path.join(UpperCamelCase       ,						'''saved_model'''       ,						'''1'''			)
                           lowercase__				      =			tf.keras.models.load_model(UpperCamelCase			)
                           lowercase__				      =			model(UpperCamelCase			)

                           if self.is_encoder_decoder:
                                  lowercase__				      =			outputs['''encoder_hidden_states''']
                                  lowercase__				      =			outputs['''encoder_attentions''']
                           else:
                                  lowercase__				      =			outputs['''hidden_states''']
                                  lowercase__				      =			outputs['''attentions''']

                           self.assertEqual(len(UpperCamelCase			)       ,						UpperCamelCase			)

                           lowercase__				      =			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 				UpperCamelCase__		(self   :   Dict			):

             '''simple docstring'''







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

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

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



      @slow
      def 				UpperCamelCase__		(self   :   List[str]			):

             '''simple docstring'''







             lowercase__      ,lowercase__				      =			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"]
             lowercase__				      =			{
                 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
                           lowercase__				      =			TFSharedEmbeddings(self.model_tester.vocab_size       ,						self.model_tester.hidden_size       ,						name='''shared'''			)
                           lowercase__				      =			False
                           lowercase__				      =			main_layer_class(UpperCamelCase       ,						embed_tokens=UpperCamelCase			)
                    else:
                           lowercase__				      =			main_layer_class(UpperCamelCase			)

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

                    if hasattr(self.model_tester       ,						'''num_labels'''			):
                           lowercase__				      =			self.model_tester.num_labels
                    else:
                           lowercase__				      =			2

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

                    lowercase__				      =			main_layer(UpperCamelCase			)[0]
                    lowercase__				      =			tf.keras.layers.Dense(UpperCamelCase       ,						activation='''softmax'''       ,						name='''outputs'''			)(UpperCamelCase			)
                    lowercase__				      =			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:
                           lowercase__				      =			os.path.join(UpperCamelCase       ,						'''keras_model.h5'''			)
                           model.save(UpperCamelCase			)
                           if "T5" in main_layer_class.__name__:
                                  lowercase__				      =			tf.keras.models.load_model(
                                      UpperCamelCase       ,						custom_objects={
                                          main_layer_class.__name__: main_layer_class,
                                          '''TFSharedEmbeddings''': TFSharedEmbeddings,
                                      }       ,						)
                           else:
                                  lowercase__				      =			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 				UpperCamelCase__		(self   :   Dict			):

             '''simple docstring'''







             lowercase__      ,lowercase__				      =			self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes[:2]:
                    lowercase__				      =			model_class(UpperCamelCase			)

                    lowercase__				      =			copy.deepcopy(UpperCamelCase			)

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

                    if not self.is_encoder_decoder:
                           lowercase__				      =			model.get_input_embeddings()(UpperCamelCase			)
                    else:
                           lowercase__				      =			model.get_input_embeddings()(UpperCamelCase			)
                           lowercase__				      =			model.get_input_embeddings()(UpperCamelCase			)

                    lowercase__				      =			self._prepare_for_class(UpperCamelCase       ,						UpperCamelCase			)

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

                    lowercase__				      =			run_in_graph_mode()
                    self.assertIsNotNone(UpperCamelCase			)



      def 				UpperCamelCase__		(self   :   Union[str, Any]       ,						UpperCamelCase   :   Dict       ,						UpperCamelCase   :   Optional[Any]			):

             '''simple docstring'''







             lowercase__				      =			[]
             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
             lowercase__				      =			[]
             while len(UpperCamelCase			) < num_bad_tokens:
                    lowercase__				      =			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 				UpperCamelCase__		(self   :   int       ,						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 				UpperCamelCase__		(self   :   Optional[int]       ,						UpperCamelCase   :   Dict       ,						UpperCamelCase   :   Tuple			):

             '''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








'''simple docstring'''

def _SCREAMING_SNAKE_CASE				(A	)	->	int:






       """simple docstring"""







       return 1 if input_a == 0 else 0

def _SCREAMING_SNAKE_CASE				()	->	None:






       """simple docstring"""







       assert not_gate(0	) == 1
       assert not_gate(1	) == 0


if __name__ == "__main__":
 print(not_gate(0))
 print(not_gate(1))




