





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:
												UpperCamelCase_     =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)]
																								)
																								UpperCamelCase_     =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		:
				UpperCamelCase			=					None
				UpperCamelCase			=					()
				UpperCamelCase			=					()
				UpperCamelCase			=					True
				UpperCamelCase			=					True
				UpperCamelCase			=					True
				UpperCamelCase			=					False






				def        snake_case			(				self	:					List[str],				lowerCAmelCase__	:					Any,				lowerCAmelCase__	:					int,				lowerCAmelCase__	:					Dict=False							)			->   dict:




										'''simple docstring'''



										_UpperCamelCase    :      List[Any]						     =					copy.deepcopy(lowerCAmelCase__							)

										if model_class in get_values(lowerCAmelCase__							):
																_UpperCamelCase    :      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__							):
																						_UpperCamelCase    :      int						     =					tf.ones(self.model_tester.batch_size,				dtype=tf.intaa							)
																elif model_class in get_values(lowerCAmelCase__							):
																						_UpperCamelCase    :      Dict						     =					tf.zeros(self.model_tester.batch_size,				dtype=tf.intaa							)
																						_UpperCamelCase    :      Any						     =					tf.zeros(self.model_tester.batch_size,				dtype=tf.intaa							)
																elif model_class in [
																    *get_values(lowerCAmelCase__							),
																    *get_values(lowerCAmelCase__							),
																]:
																						_UpperCamelCase    :      Dict						     =					tf.zeros(self.model_tester.batch_size,				dtype=tf.intaa							)
																elif model_class in get_values(lowerCAmelCase__							):
																						_UpperCamelCase    :      int						     =					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    :      Union[str, Any]						     =					tf.zeros(
																						    (self.model_tester.batch_size, self.model_tester.seq_length),				dtype=tf.intaa							)
										return inputs_dict






				@slow
				def        snake_case			(				self	:					str							)			->   Optional[Any]:




										'''simple docstring'''



										_UpperCamelCase    :      int						     =					self.model_tester.prepare_config_and_inputs_for_common()
										for model_class in self.all_model_classes[:2]:
																_UpperCamelCase    :      List[Any]						     =					self._prepare_for_class(lowerCAmelCase__,				lowerCAmelCase__							)
																_UpperCamelCase    :      List[Any]						     =					model_class(lowerCAmelCase__							)

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

																_UpperCamelCase    :      Union[str, Any]						     =					run_in_graph_mode()
																self.assertIsNotNone(lowerCAmelCase__							)






				@slow
				def        snake_case			(				self	:					List[str]							)			->   List[Any]:




										'''simple docstring'''



										_UpperCamelCase    :      str						     =					self.model_tester.prepare_config_and_inputs_for_common()
										for model_class in self.all_model_classes[:2]:
																_UpperCamelCase    :      Any						     =					self._prepare_for_class(lowerCAmelCase__,				lowerCAmelCase__							)
																_UpperCamelCase    :      List[str]						     =					model_class(lowerCAmelCase__							)

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

																_UpperCamelCase    :      Any						     =					run_in_graph_mode()
																self.assertIsNotNone(lowerCAmelCase__							)






				@slow
				def        snake_case			(				self	:					Optional[int]							)			->   Any:




										'''simple docstring'''



										_UpperCamelCase    :      Union[str, Any]						     =					self.model_tester.prepare_config_and_inputs_for_common()
										for model_class in self.all_model_classes[:2]:
																_UpperCamelCase    :      Tuple						     =					model_class(lowerCAmelCase__							)
																if getattr(lowerCAmelCase__,				'''hf_compute_loss''',				lowerCAmelCase__							):
																						# Test that model correctly compute the loss with kwargs
																						_UpperCamelCase    :      Optional[Any]						     =					self._prepare_for_class(inputs_dict.copy(),				lowerCAmelCase__,				return_labels=lowerCAmelCase__							)
																						# Is there a better way to remove these decoder inputs?
																						_UpperCamelCase    :      Any						     =					{
																						    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    :      Tuple						     =					possible_label_cols.intersection(set(lowerCAmelCase__							)							)
																						self.assertGreater(len(lowerCAmelCase__							),				0,				msg='''No matching label names found!'''							)
																						_UpperCamelCase    :      Tuple						     =					{key: val for key, val in prepared_for_class.items() if key in label_names}
																						_UpperCamelCase    :      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
																						_UpperCamelCase    :      Optional[int]						     =					model.fit(
																						    lowerCAmelCase__,				validation_data=lowerCAmelCase__,				steps_per_epoch=1,				validation_steps=1,				shuffle=lowerCAmelCase__,				verbose=0,				)
																						_UpperCamelCase    :      Tuple						     =					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    :      Optional[int]						     =					model_class(lowerCAmelCase__							)
																						model.compile(optimizer=tf.keras.optimizers.SGD(0.0							),				jit_compile=lowerCAmelCase__							)
																						_UpperCamelCase    :      Optional[Any]						     =					model.fit(
																						    lowerCAmelCase__,				lowerCAmelCase__,				validation_data=(inputs_minus_labels, labels),				steps_per_epoch=1,				validation_steps=1,				shuffle=lowerCAmelCase__,				verbose=0,				)

																						_UpperCamelCase    :      Union[str, Any]						     =					history.history['''loss'''][0]
																						self.assertTrue(not isnan(lowerCAmelCase__							)							)
																						_UpperCamelCase    :      List[Any]						     =					history.history['''val_loss'''][0]
																						self.assertTrue(not isnan(lowerCAmelCase__							)							)






				@slow
				def        snake_case			(				self	:					int							)			->   str:




										'''simple docstring'''



										_UpperCamelCase    :      Optional[Any]						     =					self.model_tester.prepare_config_and_inputs_for_common()
										_UpperCamelCase    :      List[Any]						     =					True
										_UpperCamelCase    :      Optional[int]						     =					True

										if hasattr(lowerCAmelCase__,				'''use_cache'''							):
																_UpperCamelCase    :      Optional[Any]						     =					True

										_UpperCamelCase    :      List[Any]						     =					getattr(self.model_tester,				'''encoder_seq_length''',				self.model_tester.seq_length							)
										_UpperCamelCase    :      Dict						     =					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    :      Tuple						     =					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    :      Tuple						     =					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    :      List[str]						     =					os.path.join(lowerCAmelCase__,				'''saved_model''',				'''1'''							)
																						_UpperCamelCase    :      List[str]						     =					tf.keras.models.load_model(lowerCAmelCase__							)
																						_UpperCamelCase    :      Any						     =					model(lowerCAmelCase__							)

																						if self.is_encoder_decoder:
																												_UpperCamelCase    :      str						     =					outputs['''encoder_hidden_states''']
																												_UpperCamelCase    :      Any						     =					outputs['''encoder_attentions''']
																						else:
																												_UpperCamelCase    :      Dict						     =					outputs['''hidden_states''']
																												_UpperCamelCase    :      Optional[int]						     =					outputs['''attentions''']

																						self.assertEqual(len(lowerCAmelCase__							),				lowerCAmelCase__							)

																						_UpperCamelCase    :      Tuple						     =					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        snake_case			(				self	:					Any							)			->   Any:




										'''simple docstring'''



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

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

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






				@slow
				def        snake_case			(				self	:					str							)			->   int:




										'''simple docstring'''



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

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




										'''simple docstring'''



										_UpperCamelCase    :      Dict						     =					self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes[:2]:
																_UpperCamelCase    :      Dict						     =					model_class(lowerCAmelCase__							)

																_UpperCamelCase    :      Optional[int]						     =					copy.deepcopy(lowerCAmelCase__							)

																if not self.is_encoder_decoder:
																						_UpperCamelCase    :      Union[str, Any]						     =					inputs['''input_ids''']
																						del inputs["input_ids"]
																else:
																						_UpperCamelCase    :      Tuple						     =					inputs['''input_ids''']
																						_UpperCamelCase    :      Optional[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    :      List[Any]						     =					model.get_input_embeddings()(lowerCAmelCase__							)
																						_UpperCamelCase    :      List[Any]						     =					model.get_input_embeddings()(lowerCAmelCase__							)

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

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

																_UpperCamelCase    :      str						     =					run_in_graph_mode()
																self.assertIsNotNone(lowerCAmelCase__							)






				def        snake_case			(				self	:					Tuple,				lowerCAmelCase__	:					List[str],				lowerCAmelCase__	:					Optional[Any]							)			->   Optional[int]:




										'''simple docstring'''



										_UpperCamelCase    :      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
										_UpperCamelCase    :      int						     =					[]
										while len(lowerCAmelCase__							) < num_bad_tokens:
																_UpperCamelCase    :      List[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        snake_case			(				self	:					Optional[int],				lowerCAmelCase__	:					Optional[Any]							)			->   int:




										'''simple docstring'''



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






				def        snake_case			(				self	:					Tuple,				lowerCAmelCase__	:					List[str],				lowerCAmelCase__	:					Tuple							)			->   Optional[Any]:




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










"""simple docstring"""




UpperCamelCase_     =1_000





def  a_       (							_lowercase     ):
						return x * x





def  a_       (							_lowercase ,		_lowercase ,		_lowercase ,		_lowercase = 4     ):
						assert callable(
						    _lowercase     ), F"""the function(object) passed should be callable your input : {function}"""
						assert isinstance(_lowercase ,		(float, int)     ), F"""a should be float or integer your input : {a}"""
						assert isinstance(function(_lowercase     ) ,		(float, int)     ), (
						    "the function should return integer or float return type of your function, "
						    F"""{type(_lowercase     )}"""
						)
						assert isinstance(_lowercase ,		(float, int)     ), F"""b should be float or integer your input : {b}"""
						assert (
						    isinstance(_lowercase ,		_lowercase     ) and precision > 0
						), F"""precision should be positive integer your input : {precision}"""

						# just applying the formula of simpson for approximate integration written in
						# mentioned article in first comment of this file and above this function

						_UpperCamelCase    :      str						     =					(b - a) / N_STEPS
						_UpperCamelCase    :      Dict						     =					function(_lowercase     ) + function(_lowercase     )

						for i in range(1 ,		_lowercase     ):
												_UpperCamelCase    :      List[Any]						     =					a + h * i
												result += function(_lowercase     ) * (4 if i % 2 else 2)

						result *= h / 3
						return round(_lowercase ,		_lowercase     )


if __name__ == "__main__":
						import doctest

						doctest.testmod()



