


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									=				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									=				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     __a     :
 __lowercase					:				List[Any]  	=					None
 __lowercase					:				Any  	=					()
 __lowercase					:				int  	=					()
 __lowercase					:				Union[str, Any]  	=					True
 __lowercase					:				Dict  	=					True
 __lowercase					:				List[str]  	=					True
 __lowercase					:				List[str]  	=					False
 def 						SCREAMING_SNAKE_CASE__  (  self							,					lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__=False					)   ->							dict:




        '''simple docstring'''
        lowercase__:  str         =  copy.deepcopy(lowerCAmelCase__					)

        if model_class in get_values(lowerCAmelCase__					):
               lowercase__:  Any         =  {
                   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__					):
                      lowercase__:  Optional[int]         =  tf.ones(self.model_tester.batch_size							,					dtype=tf.intaa					)
               elif model_class in get_values(lowerCAmelCase__					):
                      lowercase__:  List[Any]         =  tf.zeros(self.model_tester.batch_size							,					dtype=tf.intaa					)
                      lowercase__:  Optional[Any]         =  tf.zeros(self.model_tester.batch_size							,					dtype=tf.intaa					)
               elif model_class in [
                   *get_values(lowerCAmelCase__					),
                   *get_values(lowerCAmelCase__					),
               ]:
                      lowercase__:  List[Any]         =  tf.zeros(self.model_tester.batch_size							,					dtype=tf.intaa					)
               elif model_class in get_values(lowerCAmelCase__					):
                      lowercase__:  Tuple         =  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__					),
               ]:
                      lowercase__:  Tuple         =  tf.zeros(
                          (self.model_tester.batch_size, self.model_tester.seq_length)							,					dtype=tf.intaa					)
        return inputs_dict
 @slow
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							int:




        '''simple docstring'''
        lowercase__     ,		lowercase__:  Any         =  self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes[:2]:
               lowercase__:  Optional[Any]         =  self._prepare_for_class(lowerCAmelCase__							,					lowerCAmelCase__					)
               lowercase__:  Optional[int]         =  model_class(lowerCAmelCase__					)

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

               lowercase__:  Optional[Any]         =  run_in_graph_mode()
               self.assertIsNotNone(lowerCAmelCase__					)
 @slow
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							Optional[int]:




        '''simple docstring'''
        lowercase__     ,		lowercase__:  Any         =  self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes[:2]:
               lowercase__:  Dict         =  self._prepare_for_class(lowerCAmelCase__							,					lowerCAmelCase__					)
               lowercase__:  Dict         =  model_class(lowerCAmelCase__					)

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

               lowercase__:  List[Any]         =  run_in_graph_mode()
               self.assertIsNotNone(lowerCAmelCase__					)
 @slow
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							Tuple:




        '''simple docstring'''
        # This is a copy of the test_keras_fit method, but we use XLA compilation instead of eager
        lowercase__     ,		lowercase__:  List[Any]         =  self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes[:2]:
               lowercase__:  Optional[int]         =  model_class(lowerCAmelCase__					)
               if getattr(lowerCAmelCase__							,					'hf_compute_loss'							,					lowerCAmelCase__					):
                      # Test that model correctly compute the loss with kwargs
                      lowercase__:  Optional[Any]         =  self._prepare_for_class(inputs_dict.copy()							,					lowerCAmelCase__							,					return_labels=lowerCAmelCase__					)
                      # Is there a better way to remove these decoder inputs?
                      lowercase__:  List[str]         =  {
                          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__:  Tuple         =  {
                          'labels',
                          'label',
                          'label_ids',
                          'start_positions',
                          'start_position',
                          'end_positions',
                          'end_position',
                          'next_sentence_label',
                      }
                      lowercase__:  Tuple         =  possible_label_cols.intersection(set(lowerCAmelCase__					)					)
                      self.assertGreater(len(lowerCAmelCase__					)							,					0							,					msg='No matching label names found!'					)
                      lowercase__:  int         =  {key: val for key, val in prepared_for_class.items() if key in label_names}
                      lowercase__:  Dict         =  {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
                      lowercase__:  List[Any]         =  model.fit(
                          lowerCAmelCase__							,					validation_data=lowerCAmelCase__							,					steps_per_epoch=1							,					validation_steps=1							,					shuffle=lowerCAmelCase__							,					verbose=0							,					)
                      lowercase__:  Union[str, Any]         =  history.history['loss'][0]
                      self.assertTrue(not isnan(lowerCAmelCase__					)					)
                      lowercase__:  List[Any]         =  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.
                      lowercase__:  int         =  model_class(lowerCAmelCase__					)
                      model.compile(optimizer=tf.keras.optimizers.SGD(0.0					)							,					jit_compile=lowerCAmelCase__					)
                      lowercase__:  Any         =  model.fit(
                          lowerCAmelCase__							,					lowerCAmelCase__							,					validation_data=(inputs_minus_labels, labels)							,					steps_per_epoch=1							,					validation_steps=1							,					shuffle=lowerCAmelCase__							,					verbose=0							,					)

                      lowercase__:  List[str]         =  history.history['loss'][0]
                      self.assertTrue(not isnan(lowerCAmelCase__					)					)
                      lowercase__:  Optional[int]         =  history.history['val_loss'][0]
                      self.assertTrue(not isnan(lowerCAmelCase__					)					)
 @slow
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							Dict:




        '''simple docstring'''
        lowercase__     ,		lowercase__:  int         =  self.model_tester.prepare_config_and_inputs_for_common()
        lowercase__:  Tuple         =  True
        lowercase__:  Any         =  True

        if hasattr(lowerCAmelCase__							,					'use_cache'					):
               lowercase__:  Union[str, Any]         =  True

        lowercase__:  Union[str, Any]         =  getattr(self.model_tester							,					'encoder_seq_length'							,					self.model_tester.seq_length					)
        lowercase__:  Tuple         =  getattr(self.model_tester							,					'key_length'							,					lowerCAmelCase__					)

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

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

                      self.assertEqual(len(lowerCAmelCase__					)							,					lowerCAmelCase__					)

                      lowercase__:  Optional[Any]         =  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 						SCREAMING_SNAKE_CASE__  (  self					)   ->							Any:




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

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

                      self.assertIsNotNone(lowerCAmelCase__					)
        finally:
               tf.keras.mixed_precision.set_global_policy('float32'					)
 @slow
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							int:




        '''simple docstring'''
        lowercase__     ,		lowercase__:  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"]
        lowercase__:  Optional[Any]         =  {
            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
                      lowercase__:  Union[str, Any]         =  TFSharedEmbeddings(self.model_tester.vocab_size							,					self.model_tester.hidden_size							,					name='shared'					)
                      lowercase__:  Tuple         =  False
                      lowercase__:  List[str]         =  main_layer_class(lowerCAmelCase__							,					embed_tokens=lowerCAmelCase__					)
               else:
                      lowercase__:  Dict         =  main_layer_class(lowerCAmelCase__					)

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

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

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




        '''simple docstring'''
        lowercase__     ,		lowercase__:  str         =  self.model_tester.prepare_config_and_inputs_for_common()

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

               lowercase__:  Union[str, Any]         =  copy.deepcopy(lowerCAmelCase__					)

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

               if not self.is_encoder_decoder:
                      lowercase__:  List[str]         =  model.get_input_embeddings()(lowerCAmelCase__					)
               else:
                      lowercase__:  str         =  model.get_input_embeddings()(lowerCAmelCase__					)
                      lowercase__:  Dict         =  model.get_input_embeddings()(lowerCAmelCase__					)

               lowercase__:  str         =  self._prepare_for_class(lowerCAmelCase__							,					lowerCAmelCase__					)

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

               lowercase__:  Any         =  run_in_graph_mode()
               self.assertIsNotNone(lowerCAmelCase__					)
 def 						SCREAMING_SNAKE_CASE__  (  self							,					lowerCAmelCase__							,					lowerCAmelCase__					)   ->							List[str]:




        '''simple docstring'''
        # special tokens cannot be bad tokens
        lowercase__:  Union[str, Any]         =  []
        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__:  List[str]         =  []
        while len(lowerCAmelCase__					) < num_bad_tokens:
               lowercase__:  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(lowerCAmelCase__					)
        return bad_tokens
 def 						SCREAMING_SNAKE_CASE__  (  self							,					lowerCAmelCase__					)   ->							Optional[Any]:




        '''simple docstring'''
        for token_id in output_ids[0].numpy().tolist():
               self.assertGreaterEqual(lowerCAmelCase__							,					0					)
               self.assertLess(lowerCAmelCase__							,					self.model_tester.vocab_size					)
 def 						SCREAMING_SNAKE_CASE__  (  self							,					lowerCAmelCase__							,					lowerCAmelCase__					)   ->							Union[str, Any]:




        '''simple docstring'''
        # for all bad word tokens
        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






from typing import TYPE_CHECKING

from ...utils import _LazyModule


__lowerCAmelCase									=				{'''tokenization_phobert''': ['''PhobertTokenizer''']}


if TYPE_CHECKING:
 from .tokenization_phobert import PhobertTokenizer

else:
 import sys

 __lowerCAmelCase									=				_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)


