
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	:      List[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)]
								)
								__SCREAMING_SNAKE_CASE	:      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 						__magic_name__      :
		_lowerCAmelCase						=		None
		_lowerCAmelCase						=		()
		_lowerCAmelCase						=		()
		_lowerCAmelCase						=		True
		_lowerCAmelCase						=		True
		_lowerCAmelCase						=		True
		_lowerCAmelCase						=		False





		def 					_A							(							self			:		Optional[Any]  ,     lowerCamelCase__			:		Tuple  ,     lowerCamelCase__			:		Dict  ,     lowerCamelCase__			:		Optional[Any]=False ):
							lowerCAmelCase						:    List[Any]     		=   copy.deepcopy(_snake_case )

							if model_class in get_values(_snake_case ):
												lowerCAmelCase						:    Union[str, Any]     		=   {
												    k: tf.tile(tf.expand_dims(_snake_case  ,     1 )  ,     (1, self.model_tester.num_choices) + (1,) * (v.ndim - 1) )
												    if isinstance(_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(_snake_case ):
																	lowerCAmelCase						:    Union[str, Any]     		=   tf.ones(self.model_tester.batch_size  ,     dtype=tf.intaa )
												elif model_class in get_values(_snake_case ):
																	lowerCAmelCase						:    List[Any]     		=   tf.zeros(self.model_tester.batch_size  ,     dtype=tf.intaa )
																	lowerCAmelCase						:    List[Any]     		=   tf.zeros(self.model_tester.batch_size  ,     dtype=tf.intaa )
												elif model_class in [
												    *get_values(_snake_case ),
												    *get_values(_snake_case ),
												]:
																	lowerCAmelCase						:    Optional[Any]     		=   tf.zeros(self.model_tester.batch_size  ,     dtype=tf.intaa )
												elif model_class in get_values(_snake_case ):
																	lowerCAmelCase						:    Optional[Any]     		=   tf.zeros(self.model_tester.batch_size  ,     dtype=tf.intaa )
												elif model_class in [
												    *get_values(_snake_case ),
												    *get_values(_snake_case ),
												    *get_values(_snake_case ),
												    *get_values(_snake_case ),
												    *get_values(_snake_case ),
												]:
																	lowerCAmelCase						:    Dict     		=   tf.zeros(
																	    (self.model_tester.batch_size, self.model_tester.seq_length)  ,     dtype=tf.intaa )
							return inputs_dict





		@slow
		def 					_A							(							self			:		List[Any] ):
							lowerCAmelCase						:    List[str]     		=   self.model_tester.prepare_config_and_inputs_for_common()
							for model_class in self.all_model_classes[:2]:
												lowerCAmelCase						:    Optional[int]     		=   self._prepare_for_class(_snake_case  ,     _snake_case )
												lowerCAmelCase						:    str     		=   model_class(_snake_case )

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

												lowerCAmelCase						:    int     		=   run_in_graph_mode()
												self.assertIsNotNone(_snake_case )





		@slow
		def 					_A							(							self			:		Dict ):
							lowerCAmelCase						:    Dict     		=   self.model_tester.prepare_config_and_inputs_for_common()
							for model_class in self.all_model_classes[:2]:
												lowerCAmelCase						:    str     		=   self._prepare_for_class(_snake_case  ,     _snake_case )
												lowerCAmelCase						:    List[Any]     		=   model_class(_snake_case )

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

												lowerCAmelCase						:    Union[str, Any]     		=   run_in_graph_mode()
												self.assertIsNotNone(_snake_case )





		@slow
		def 					_A							(							self			:		Union[str, Any] ):
							# This is a copy of the test_keras_fit method, but we use XLA compilation instead of eager
							lowerCAmelCase						:    Optional[int]     		=   self.model_tester.prepare_config_and_inputs_for_common()
							for model_class in self.all_model_classes[:2]:
												lowerCAmelCase						:    List[str]     		=   model_class(_snake_case )
												if getattr(_snake_case  ,     '''hf_compute_loss'''  ,     _snake_case ):
																	# Test that model correctly compute the loss with kwargs
																	lowerCAmelCase						:    Any     		=   self._prepare_for_class(inputs_dict.copy()  ,     _snake_case  ,     return_labels=_snake_case )
																	# Is there a better way to remove these decoder inputs?
																	lowerCAmelCase						:    Optional[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")
																	}

																	lowerCAmelCase						:    Any     		=   {
																	    "labels",
																	    "label",
																	    "label_ids",
																	    "start_positions",
																	    "start_position",
																	    "end_positions",
																	    "end_position",
																	    "next_sentence_label",
																	}
																	lowerCAmelCase						:    int     		=   possible_label_cols.intersection(set(_snake_case ) )
																	self.assertGreater(len(_snake_case )  ,     0  ,     msg='''No matching label names found!''' )
																	lowerCAmelCase						:    Optional[Any]     		=   {key: val for key, val in prepared_for_class.items() if key in label_names}
																	lowerCAmelCase						:    Optional[Any]     		=   {key: val for key, val in prepared_for_class.items() if key not in label_names}
																	self.assertGreater(len(_snake_case )  ,     0 )

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

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

																	lowerCAmelCase						:    str     		=   history.history["loss"][0]
																	self.assertTrue(not isnan(_snake_case ) )
																	lowerCAmelCase						:    Union[str, Any]     		=   history.history["val_loss"][0]
																	self.assertTrue(not isnan(_snake_case ) )





		@slow
		def 					_A							(							self			:		Optional[int] ):
							lowerCAmelCase						:    Dict     		=   self.model_tester.prepare_config_and_inputs_for_common()
							lowerCAmelCase						:    Tuple     		=   True
							lowerCAmelCase						:    Optional[Any]     		=   True

							if hasattr(_snake_case  ,     '''use_cache''' ):
												lowerCAmelCase						:    Union[str, Any]     		=   True

							lowerCAmelCase						:    Tuple     		=   getattr(self.model_tester  ,     '''encoder_seq_length'''  ,     self.model_tester.seq_length )
							lowerCAmelCase						:    List[Any]     		=   getattr(self.model_tester  ,     '''key_length'''  ,     _snake_case )

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

																	if self.is_encoder_decoder:
																						lowerCAmelCase						:    int     		=   outputs["encoder_hidden_states"]
																						lowerCAmelCase						:    int     		=   outputs["encoder_attentions"]
																	else:
																						lowerCAmelCase						:    Dict     		=   outputs["hidden_states"]
																						lowerCAmelCase						:    List[str]     		=   outputs["attentions"]

																	self.assertEqual(len(_snake_case )  ,     _snake_case )

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

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

																	self.assertEqual(len(_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 					_A							(							self			:		List[str] ):
							tf.keras.mixed_precision.set_global_policy('''mixed_float16''' )

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

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





		@slow
		def 					_A							(							self			:		List[str] ):
							lowerCAmelCase						:    Any     		=   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"]
							lowerCAmelCase						:    Tuple     		=   {
							    module_member
							    for model_class in self.all_model_classes
							    for module in (import_module(model_class.__module__ ),)
							    for module_member_name in dir(_snake_case )
							    if module_member_name.endswith('''MainLayer''' )
							    for module_member in (getattr(_snake_case  ,     _snake_case ),)
							    if isinstance(_snake_case  ,     _snake_case )
							    and tf.keras.layers.Layer in module_member.__bases__
							    and getattr(_snake_case  ,     '''_keras_serializable'''  ,     _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
																	lowerCAmelCase						:    Optional[int]     		=   TFSharedEmbeddings(self.model_tester.vocab_size  ,     self.model_tester.hidden_size  ,     name='''shared''' )
																	lowerCAmelCase						:    str     		=   False
																	lowerCAmelCase						:    Dict     		=   main_layer_class(_snake_case  ,     embed_tokens=_snake_case )
												else:
																	lowerCAmelCase						:    List[Any]     		=   main_layer_class(_snake_case )

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

												if hasattr(self.model_tester  ,     '''num_labels''' ):
																	lowerCAmelCase						:    Union[str, Any]     		=   self.model_tester.num_labels
												else:
																	lowerCAmelCase						:    List[str]     		=   2

												lowerCAmelCase						:    Optional[int]     		=   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 )

												lowerCAmelCase						:    Dict     		=   main_layer(_snake_case )[0]
												lowerCAmelCase						:    Any     		=   tf.keras.layers.Dense(_snake_case  ,     activation='''softmax'''  ,     name='''outputs''' )(_snake_case )
												lowerCAmelCase						:    str     		=   tf.keras.models.Model(inputs=_snake_case  ,     outputs=[outputs] )

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

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





		@slow
		def 					_A							(							self			:		str ):
							lowerCAmelCase						:    List[Any]     		=   self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes[:2]:
												lowerCAmelCase						:    List[str]     		=   model_class(_snake_case )

												lowerCAmelCase						:    Union[str, Any]     		=   copy.deepcopy(_snake_case )

												if not self.is_encoder_decoder:
																	lowerCAmelCase						:    Union[str, Any]     		=   inputs["input_ids"]
																	del inputs["input_ids"]
												else:
																	lowerCAmelCase						:    int     		=   inputs["input_ids"]
																	lowerCAmelCase						:    List[str]     		=   inputs.get('''decoder_input_ids'''  ,     _snake_case )
																	del inputs["input_ids"]
																	inputs.pop('''decoder_input_ids'''  ,     _snake_case )

												if not self.is_encoder_decoder:
																	lowerCAmelCase						:    List[str]     		=   model.get_input_embeddings()(_snake_case )
												else:
																	lowerCAmelCase						:    int     		=   model.get_input_embeddings()(_snake_case )
																	lowerCAmelCase						:    Union[str, Any]     		=   model.get_input_embeddings()(_snake_case )

												lowerCAmelCase						:    Any     		=   self._prepare_for_class(_snake_case  ,     _snake_case )

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

												lowerCAmelCase						:    int     		=   run_in_graph_mode()
												self.assertIsNotNone(_snake_case )





		def 					_A							(							self			:		Union[str, Any]  ,     lowerCamelCase__			:		List[Any]  ,     lowerCamelCase__			:		Tuple ):
							# special tokens cannot be bad tokens
							lowerCAmelCase						:    List[str]     		=   []
							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
							lowerCAmelCase						:    List[str]     		=   []
							while len(_snake_case ) < num_bad_tokens:
												lowerCAmelCase						:    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(_snake_case )
							return bad_tokens





		def 					_A							(							self			:		Union[str, Any]  ,     lowerCamelCase__			:		Any ):
							for token_id in output_ids[0].numpy().tolist():
												self.assertGreaterEqual(_snake_case  ,     0 )
												self.assertLess(_snake_case  ,     self.model_tester.vocab_size )





		def 					_A							(							self			:		int  ,     lowerCamelCase__			:		List[str]  ,     lowerCamelCase__			:		Dict ):
							# 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(_snake_case )  ,     len(_snake_case ) ):
																						# if tokens match
																						if generated_ids_slice[i - len(_snake_case ) : i] == bad_word_ids:
																											return True
							return False



'''simple docstring'''
import os
import unittest

import pytest
from parameterized import parameterized

from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device


# skipping in unittest tests

_lowerCamelCase	 =			[(1,)]


def       a__    (  )			-> str:


	"""simple docstring"""





	UpperCAmelCase_			:	Optional[int]      =   bool(os.getenv("RUN_SLOW"     ,							0							)							)
	if run_slow:
		assert True
	else:
		assert False, "should have been skipped"


def       a__    (  )			-> str:


	"""simple docstring"""





	UpperCAmelCase_			:	Optional[Any]      =   bool(os.getenv("RUN_SLOW"     ,							0							)							)
	if run_slow and torch_device == "cuda":
		assert True
	else:
		assert False, "should have been skipped"
@require_torch
class  _snake_case						(unittest.TestCase):






			@slow
			@require_torch_gpu
			def 						UpperCamelCase__				(   self			):
				check_slow_torch_cuda()






			@require_torch_gpu
			@slow
			def 						UpperCamelCase__				(   self			):
				check_slow_torch_cuda()






			@parameterized.expand(_snake_case			)
			@slow
			def 						UpperCamelCase__				(   self       ,_snake_case=None			):
				check_slow()




@slow
@require_torch_gpu
def       a__    (  )			-> List[Any]:


	"""simple docstring"""





	check_slow_torch_cuda()


@require_torch_gpu
@slow
def       a__    (  )			-> List[str]:


	"""simple docstring"""





	check_slow_torch_cuda()


@slow
@pytest.mark.parametrize("param"     ,							[1]							)
def       a__    (  _SCREAMING_SNAKE_CASE					:    int							)			-> List[Any]:


	"""simple docstring"""





	check_slow()





@pytest.mark.parametrize("param"     ,							[1]							)
@slow
def       a__    (  _SCREAMING_SNAKE_CASE					:    List[str]							)			-> Any:


	"""simple docstring"""





	check_slow()




