


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__						:str         =							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__						:int         =							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   :
		_a			:						Tuple    			=			None
		_a			:						int    			=			()
		_a			:						Dict    			=			()
		_a			:						Optional[int]    			=			True
		_a			:						Union[str, Any]    			=			True
		_a			:						Optional[Any]    			=			True
		_a			:						Any    			=			False
		def 			UpperCAmelCase__       (    self    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE=False		)				-> dict:




									"""simple docstring"""





									_UpperCAmelCase					=    copy.deepcopy(_SCREAMING_SNAKE_CASE		)

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

									if return_labels:
																if model_class in get_values(_SCREAMING_SNAKE_CASE		):
																							_UpperCAmelCase					=    tf.ones(self.model_tester.batch_size    ,					dtype=tf.intaa		)
																elif model_class in get_values(_SCREAMING_SNAKE_CASE		):
																							_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(_SCREAMING_SNAKE_CASE		),
																    *get_values(_SCREAMING_SNAKE_CASE		),
																]:
																							_UpperCAmelCase					=    tf.zeros(self.model_tester.batch_size    ,					dtype=tf.intaa		)
																elif model_class in get_values(_SCREAMING_SNAKE_CASE		):
																							_UpperCAmelCase					=    tf.zeros(self.model_tester.batch_size    ,					dtype=tf.intaa		)
																elif model_class in [
																    *get_values(_SCREAMING_SNAKE_CASE		),
																    *get_values(_SCREAMING_SNAKE_CASE		),
																    *get_values(_SCREAMING_SNAKE_CASE		),
																    *get_values(_SCREAMING_SNAKE_CASE		),
																    *get_values(_SCREAMING_SNAKE_CASE		),
																]:
																							_UpperCAmelCase					=    tf.zeros(
																							    (self.model_tester.batch_size, self.model_tester.seq_length)    ,					dtype=tf.intaa		)
									return inputs_dict
		@slow
		def 			UpperCAmelCase__       (    self		)				-> 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(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase					=    model_class(_SCREAMING_SNAKE_CASE		)

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

																_UpperCAmelCase					=    run_in_graph_mode()
																self.assertIsNotNone(_SCREAMING_SNAKE_CASE		)
		@slow
		def 			UpperCAmelCase__       (    self		)				-> Union[str, 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(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase					=    model_class(_SCREAMING_SNAKE_CASE		)

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

																_UpperCAmelCase					=    run_in_graph_mode()
																self.assertIsNotNone(_SCREAMING_SNAKE_CASE		)
		@slow
		def 			UpperCAmelCase__       (    self		)				-> List[Any]:




									"""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(_SCREAMING_SNAKE_CASE		)
																if getattr(_SCREAMING_SNAKE_CASE    ,					'hf_compute_loss'    ,					_SCREAMING_SNAKE_CASE		):
																							# Test that model correctly compute the loss with kwargs
																							_UpperCAmelCase					=    self._prepare_for_class(inputs_dict.copy()    ,					_SCREAMING_SNAKE_CASE    ,					return_labels=_SCREAMING_SNAKE_CASE		)
																							# 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(_SCREAMING_SNAKE_CASE		)		)
																							self.assertGreater(len(_SCREAMING_SNAKE_CASE		)    ,					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(_SCREAMING_SNAKE_CASE		)    ,					0		)

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

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

																							_UpperCAmelCase					=    history.history['loss'][0]
																							self.assertTrue(not isnan(_SCREAMING_SNAKE_CASE		)		)
																							_UpperCAmelCase					=    history.history['val_loss'][0]
																							self.assertTrue(not isnan(_SCREAMING_SNAKE_CASE		)		)
		@slow
		def 			UpperCAmelCase__       (    self		)				-> str:




									"""simple docstring"""





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

									if hasattr(_SCREAMING_SNAKE_CASE    ,					'use_cache'		):
																_UpperCAmelCase					=    True

									_UpperCAmelCase					=    getattr(self.model_tester    ,					'encoder_seq_length'    ,					self.model_tester.seq_length		)
									_UpperCAmelCase					=    getattr(self.model_tester    ,					'key_length'    ,					_SCREAMING_SNAKE_CASE		)

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

																for key in list(class_inputs_dict.keys()		):
																							# Remove keys not in the serving signature, as the SavedModel will not be compiled to deal with them
																							if key not in model.input_signature:
																														del class_inputs_dict[key]
																							# Check it's a tensor, in case the inputs dict has some bools in it too
																							elif isinstance(class_inputs_dict[key]    ,					tf.Tensor		) and class_inputs_dict[key].dtype.is_integer:
																														_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(_SCREAMING_SNAKE_CASE    ,					saved_model=_SCREAMING_SNAKE_CASE		)
																							_UpperCAmelCase					=    os.path.join(_SCREAMING_SNAKE_CASE    ,					'saved_model'    ,					'1'		)
																							_UpperCAmelCase					=    tf.keras.models.load_model(_SCREAMING_SNAKE_CASE		)
																							_UpperCAmelCase					=    model(_SCREAMING_SNAKE_CASE		)

																							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(_SCREAMING_SNAKE_CASE		)    ,					_SCREAMING_SNAKE_CASE		)

																							_UpperCAmelCase					=    getattr(
																							    self.model_tester    ,					'expected_num_hidden_layers'    ,					self.model_tester.num_hidden_layers + 1		)

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

																							self.assertEqual(len(_SCREAMING_SNAKE_CASE		)    ,					self.model_tester.num_hidden_layers		)
																							self.assertListEqual(
																							    list(output_attentions[0].shape[-3:]		)    ,					[self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length]    ,					)
		@slow
		def 			UpperCAmelCase__       (    self		)				-> Union[str, Any]:




									"""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(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		)
																							_UpperCAmelCase					=    model_class(_SCREAMING_SNAKE_CASE		)
																							_UpperCAmelCase					=    model(_SCREAMING_SNAKE_CASE		)

																							self.assertIsNotNone(_SCREAMING_SNAKE_CASE		)
									finally:
																tf.keras.mixed_precision.set_global_policy('float32'		)
		@slow
		def 			UpperCAmelCase__       (    self		)				-> List[str]:




									"""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(_SCREAMING_SNAKE_CASE		)
									    if module_member_name.endswith('MainLayer'		)
									    for module_member in (getattr(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		),)
									    if isinstance(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		)
									    and tf.keras.layers.Layer in module_member.__bases__
									    and getattr(_SCREAMING_SNAKE_CASE    ,					'_keras_serializable'    ,					_SCREAMING_SNAKE_CASE		)
									}

									for main_layer_class in tf_main_layer_classes:
																# T5MainLayer needs an embed_tokens parameter when called without the inputs_embeds parameter
																if "T5" in main_layer_class.__name__:
																							# Take the same values than in TFT5ModelTester for this shared layer
																							_UpperCAmelCase					=    TFSharedEmbeddings(self.model_tester.vocab_size    ,					self.model_tester.hidden_size    ,					name='shared'		)
																							_UpperCAmelCase					=    False
																							_UpperCAmelCase					=    main_layer_class(_SCREAMING_SNAKE_CASE    ,					embed_tokens=_SCREAMING_SNAKE_CASE		)
																else:
																							_UpperCAmelCase					=    main_layer_class(_SCREAMING_SNAKE_CASE		)

																_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(_SCREAMING_SNAKE_CASE		)[0]
																_UpperCAmelCase					=    tf.keras.layers.Dense(_SCREAMING_SNAKE_CASE    ,					activation='softmax'    ,					name='outputs'		)(_SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase					=    tf.keras.models.Model(inputs=_SCREAMING_SNAKE_CASE    ,					outputs=[outputs]		)

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

																with tempfile.TemporaryDirectory() as tmpdirname:
																							_UpperCAmelCase					=    os.path.join(_SCREAMING_SNAKE_CASE    ,					'keras_model.h5'		)
																							model.save(_SCREAMING_SNAKE_CASE		)
																							if "T5" in main_layer_class.__name__:
																														_UpperCAmelCase					=    tf.keras.models.load_model(
																														    _SCREAMING_SNAKE_CASE    ,					custom_objects={
																														        main_layer_class.__name__: main_layer_class,
																														        'TFSharedEmbeddings': TFSharedEmbeddings,
																														    }    ,					)
																							else:
																														_UpperCAmelCase					=    tf.keras.models.load_model(
																														    _SCREAMING_SNAKE_CASE    ,					custom_objects={main_layer_class.__name__: main_layer_class}		)
																							assert isinstance(_SCREAMING_SNAKE_CASE    ,					tf.keras.Model		)
																							model(_SCREAMING_SNAKE_CASE		)
		@slow
		def 			UpperCAmelCase__       (    self		)				-> Union[str, Any]:




									"""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(_SCREAMING_SNAKE_CASE		)

																_UpperCAmelCase					=    copy.deepcopy(_SCREAMING_SNAKE_CASE		)

																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'    ,					_SCREAMING_SNAKE_CASE		)
																							del inputs["input_ids"]
																							inputs.pop('decoder_input_ids'    ,					_SCREAMING_SNAKE_CASE		)

																if not self.is_encoder_decoder:
																							_UpperCAmelCase					=    model.get_input_embeddings()(_SCREAMING_SNAKE_CASE		)
																else:
																							_UpperCAmelCase					=    model.get_input_embeddings()(_SCREAMING_SNAKE_CASE		)
																							_UpperCAmelCase					=    model.get_input_embeddings()(_SCREAMING_SNAKE_CASE		)

																_UpperCAmelCase					=    self._prepare_for_class(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		)

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

																_UpperCAmelCase					=    run_in_graph_mode()
																self.assertIsNotNone(_SCREAMING_SNAKE_CASE		)
		def 			UpperCAmelCase__       (    self    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		)				-> Union[str, Any]:




									"""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(_SCREAMING_SNAKE_CASE		) < 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(_SCREAMING_SNAKE_CASE		)
									return bad_tokens
		def 			UpperCAmelCase__       (    self    ,					_SCREAMING_SNAKE_CASE		)				-> List[str]:




									"""simple docstring"""





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






		def 			UpperCAmelCase__       (    self    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		)				-> Optional[int]:




									"""simple docstring"""





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




import warnings

from .generation import GenerationMixin

class 	__a   (       UpperCAmelCase    ):
		# warning at import time
		warnings.warn(
		    'Importing `GenerationMixin` from `src/transformers/generation_utils.py` is deprecated and will '
		    'be removed in Transformers v5. Import as `from transformers import GenerationMixin` instead.'					,   UpperCAmelCase					,   )
