


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:
    _SCREAMING_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)]
        )
        _SCREAMING_SNAKE_CASE :					List[str]			=	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  UpperCamelCase__ :
      a__   :     int								=						None
      a__   :     Union[str, Any]								=						()
      a__   :     List[str]								=						()
      a__   :     Optional[Any]								=						True
      a__   :     Optional[Any]								=						True
      a__   :     Any								=						True
      a__   :     int								=						False






      def 	__lowercase(     self					:   Optional[int],		__lowerCamelCase					:   Union[str, Any],		__lowerCamelCase					:   int,		__lowerCamelCase					:   int=False   )  -> dict:
           UpperCamelCase__							:  Any								=     copy.deepcopy(__lowerCamelCase   )

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

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






      @slow
      def 	__lowercase(     self					:   Any   )  -> Any:
           UpperCamelCase__							:  List[str]								=     self.model_tester.prepare_config_and_inputs_for_common()
           for model_class in self.all_model_classes[:2]:
                UpperCamelCase__							:  str								=     self._prepare_for_class(__lowerCamelCase,		__lowerCamelCase   )
                UpperCamelCase__							:  List[str]								=     model_class(__lowerCamelCase   )

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

                UpperCamelCase__							:  int								=     run_in_graph_mode()
                self.assertIsNotNone(__lowerCamelCase   )






      @slow
      def 	__lowercase(     self					:   Tuple   )  -> int:
           UpperCamelCase__							:  List[str]								=     self.model_tester.prepare_config_and_inputs_for_common()
           for model_class in self.all_model_classes[:2]:
                UpperCamelCase__							:  str								=     self._prepare_for_class(__lowerCamelCase,		__lowerCamelCase   )
                UpperCamelCase__							:  Union[str, Any]								=     model_class(__lowerCamelCase   )

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

                UpperCamelCase__							:  int								=     run_in_graph_mode()
                self.assertIsNotNone(__lowerCamelCase   )






      @slow
      def 	__lowercase(     self					:   Union[str, Any]   )  -> List[Any]:
           UpperCamelCase__							:  Optional[int]								=     self.model_tester.prepare_config_and_inputs_for_common()
           for model_class in self.all_model_classes[:2]:
                UpperCamelCase__							:  List[Any]								=     model_class(__lowerCamelCase   )
                if getattr(__lowerCamelCase,		'''hf_compute_loss''',		__lowerCamelCase   ):
                     # Test that model correctly compute the loss with kwargs
                     UpperCamelCase__							:  List[Any]								=     self._prepare_for_class(inputs_dict.copy(),		__lowerCamelCase,		return_labels=__lowerCamelCase   )
                     # Is there a better way to remove these decoder inputs?
                     UpperCamelCase__							:  int								=     {
                         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__							:  Dict								=     {
                         "labels",
                         "label",
                         "label_ids",
                         "start_positions",
                         "start_position",
                         "end_positions",
                         "end_position",
                         "next_sentence_label",
                     }
                     UpperCamelCase__							:  str								=     possible_label_cols.intersection(set(__lowerCamelCase   )   )
                     self.assertGreater(len(__lowerCamelCase   ),		0,		msg='''No matching label names found!'''   )
                     UpperCamelCase__							:  Optional[Any]								=     {key: val for key, val in prepared_for_class.items() if key in label_names}
                     UpperCamelCase__							:  int								=     {key: val for key, val in prepared_for_class.items() if key not in label_names}
                     self.assertGreater(len(__lowerCamelCase   ),		0   )

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

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

                     UpperCamelCase__							:  str								=     history.history["loss"][0]
                     self.assertTrue(not isnan(__lowerCamelCase   )   )
                     UpperCamelCase__							:  List[Any]								=     history.history["val_loss"][0]
                     self.assertTrue(not isnan(__lowerCamelCase   )   )






      @slow
      def 	__lowercase(     self					:   Tuple   )  -> Any:
           UpperCamelCase__							:  Union[str, Any]								=     self.model_tester.prepare_config_and_inputs_for_common()
           UpperCamelCase__							:  Dict								=     True
           UpperCamelCase__							:  Any								=     True

           if hasattr(__lowerCamelCase,		'''use_cache'''   ):
                UpperCamelCase__							:  Any								=     True

           UpperCamelCase__							:  int								=     getattr(self.model_tester,		'''encoder_seq_length''',		self.model_tester.seq_length   )
           UpperCamelCase__							:  Tuple								=     getattr(self.model_tester,		'''key_length''',		__lowerCamelCase   )

           for model_class in self.all_model_classes[:2]:
                UpperCamelCase__							:  Dict								=     self._prepare_for_class(__lowerCamelCase,		__lowerCamelCase   )
                UpperCamelCase__							:  int								=     model_class(__lowerCamelCase   )
                model.build()
                UpperCamelCase__							:  str								=     len(model(__lowerCamelCase   )   )

                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__							:  List[str]								=     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(__lowerCamelCase,		saved_model=__lowerCamelCase   )
                     UpperCamelCase__							:  str								=     os.path.join(__lowerCamelCase,		'''saved_model''',		'''1'''   )
                     UpperCamelCase__							:  Tuple								=     tf.keras.models.load_model(__lowerCamelCase   )
                     UpperCamelCase__							:  Tuple								=     model(__lowerCamelCase   )

                     if self.is_encoder_decoder:
                          UpperCamelCase__							:  int								=     outputs["encoder_hidden_states"]
                          UpperCamelCase__							:  str								=     outputs["encoder_attentions"]
                     else:
                          UpperCamelCase__							:  int								=     outputs["hidden_states"]
                          UpperCamelCase__							:  Optional[Any]								=     outputs["attentions"]

                     self.assertEqual(len(__lowerCamelCase   ),		__lowerCamelCase   )

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

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

                     self.assertEqual(len(__lowerCamelCase   ),		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 	__lowercase(     self					:   List[Any]   )  -> Union[str, Any]:
           tf.keras.mixed_precision.set_global_policy('''mixed_float16'''   )

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

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






      @slow
      def 	__lowercase(     self					:   int   )  -> Union[str, Any]:
           UpperCamelCase__							:  Dict								=     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__							:  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(__lowerCamelCase   )
               if module_member_name.endswith('''MainLayer'''   )
               for module_member in (getattr(__lowerCamelCase,		__lowerCamelCase   ),)
               if isinstance(__lowerCamelCase,		__lowerCamelCase   )
               and tf.keras.layers.Layer in module_member.__bases__
               and getattr(__lowerCamelCase,		'''_keras_serializable''',		__lowerCamelCase   )
           }

           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__							:  Tuple								=     TFSharedEmbeddings(self.model_tester.vocab_size,		self.model_tester.hidden_size,		name='''shared'''   )
                     UpperCamelCase__							:  List[str]								=     False
                     UpperCamelCase__							:  List[Any]								=     main_layer_class(__lowerCamelCase,		embed_tokens=__lowerCamelCase   )
                else:
                     UpperCamelCase__							:  Any								=     main_layer_class(__lowerCamelCase   )

                UpperCamelCase__							:  str								=     {
                    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__							:  List[Any]								=     self.model_tester.num_labels
                else:
                     UpperCamelCase__							:  Tuple								=     2

                UpperCamelCase__							:  List[str]								=     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__							:  Any								=     main_layer(__lowerCamelCase   )[0]
                UpperCamelCase__							:  str								=     tf.keras.layers.Dense(__lowerCamelCase,		activation='''softmax''',		name='''outputs'''   )(__lowerCamelCase   )
                UpperCamelCase__							:  str								=     tf.keras.models.Model(inputs=__lowerCamelCase,		outputs=[outputs]   )

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

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






      @slow
      def 	__lowercase(     self					:   str   )  -> int:
           UpperCamelCase__							:  int								=     self.model_tester.prepare_config_and_inputs_for_common()

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

                UpperCamelCase__							:  str								=     copy.deepcopy(__lowerCamelCase   )

                if not self.is_encoder_decoder:
                     UpperCamelCase__							:  int								=     inputs["input_ids"]
                     del inputs["input_ids"]
                else:
                     UpperCamelCase__							:  Dict								=     inputs["input_ids"]
                     UpperCamelCase__							:  Any								=     inputs.get('''decoder_input_ids''',		__lowerCamelCase   )
                     del inputs["input_ids"]
                     inputs.pop('''decoder_input_ids''',		__lowerCamelCase   )

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

                UpperCamelCase__							:  Dict								=     self._prepare_for_class(__lowerCamelCase,		__lowerCamelCase   )

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

                UpperCamelCase__							:  int								=     run_in_graph_mode()
                self.assertIsNotNone(__lowerCamelCase   )






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






      def 	__lowercase(     self					:   List[str],		__lowerCamelCase					:   int   )  -> Any:






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






      def 	__lowercase(     self					:   Union[str, Any],		__lowerCamelCase					:   Tuple,		__lowerCamelCase					:   Optional[Any]   )  -> List[str]:






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








import warnings

from ...utils import logging
from .image_processing_mobilenet_va import MobileNetVaImageProcessor


_SCREAMING_SNAKE_CASE :					List[str]			=	logging.get_logger(__name__)

class  UpperCamelCase__ ( __lowerCamelCase			):






      def __init__(     self					:   List[str],		*__lowerCamelCase					:   Optional[int],		**__lowerCamelCase					:   Dict   )  -> None:
           warnings.warn(
               '''The class MobileNetV1FeatureExtractor is deprecated and will be removed in version 5 of Transformers.'''
               ''' Please use MobileNetV1ImageProcessor instead.''',		__lowerCamelCase,		)
           super().__init__(*__lowerCamelCase,		**__lowerCamelCase   )




