


'''simple docstring'''


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 						lowerCAmelCase__      :
			lowerCAmelCase			:						Optional[Any]							 =  None
			lowerCAmelCase			:						Tuple							 =  ()
			lowerCAmelCase			:						Any							 =  ()
			lowerCAmelCase			:						Union[str, Any]							 =  True
			lowerCAmelCase			:						Any							 =  True
			lowerCAmelCase			:						int							 =  True
			lowerCAmelCase			:						List[Any]							 =  False




			def 				lowerCAmelCase__     (		self	:  Union[str, Any]						,	lowerCamelCase__	:  Tuple						,	lowerCamelCase__	:  Tuple						,	lowerCamelCase__	:  Dict=False			) ->dict:



					'''simple docstring'''







					_UpperCAmelCase		:							Optional[int]             =   copy.deepcopy(_UpperCAmelCase			)

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

					if return_labels:
							if model_class in get_values(_UpperCAmelCase			):
									_UpperCAmelCase		:							Any             =   tf.ones(self.model_tester.batch_size						,	dtype=tf.intaa			)
							elif model_class in get_values(_UpperCAmelCase			):
									_UpperCAmelCase		:							Optional[Any]             =   tf.zeros(self.model_tester.batch_size						,	dtype=tf.intaa			)
									_UpperCAmelCase		:							Dict             =   tf.zeros(self.model_tester.batch_size						,	dtype=tf.intaa			)
							elif model_class in [
							    *get_values(_UpperCAmelCase			),
							    *get_values(_UpperCAmelCase			),
							]:
									_UpperCAmelCase		:							Any             =   tf.zeros(self.model_tester.batch_size						,	dtype=tf.intaa			)
							elif model_class in get_values(_UpperCAmelCase			):
									_UpperCAmelCase		:							List[Any]             =   tf.zeros(self.model_tester.batch_size						,	dtype=tf.intaa			)
							elif model_class in [
							    *get_values(_UpperCAmelCase			),
							    *get_values(_UpperCAmelCase			),
							    *get_values(_UpperCAmelCase			),
							    *get_values(_UpperCAmelCase			),
							    *get_values(_UpperCAmelCase			),
							]:
									_UpperCAmelCase		:							Tuple             =   tf.zeros(
									    (self.model_tester.batch_size, self.model_tester.seq_length)						,	dtype=tf.intaa			)
					return inputs_dict




			@slow
			def 				lowerCAmelCase__     (		self	:  Tuple			) ->List[str]:



					'''simple docstring'''







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

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

							_UpperCAmelCase		:							Optional[Any]             =   run_in_graph_mode()
							self.assertIsNotNone(_UpperCAmelCase			)




			@slow
			def 				lowerCAmelCase__     (		self	:  List[str]			) ->int:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							Any             =   self.model_tester.prepare_config_and_inputs_for_common()
					for model_class in self.all_model_classes[:2]:
							_UpperCAmelCase		:							Tuple             =   self._prepare_for_class(_UpperCAmelCase						,	_UpperCAmelCase			)
							_UpperCAmelCase		:							Dict             =   model_class(_UpperCAmelCase			)

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

							_UpperCAmelCase		:							Tuple             =   run_in_graph_mode()
							self.assertIsNotNone(_UpperCAmelCase			)




			@slow
			def 				lowerCAmelCase__     (		self	:  Optional[int]			) ->List[str]:



					'''simple docstring'''







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

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

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

									_UpperCAmelCase		:							Optional[Any]             =   history.history["loss"][0]
									self.assertTrue(not isnan(_UpperCAmelCase			)			)
									_UpperCAmelCase		:							str             =   history.history["val_loss"][0]
									self.assertTrue(not isnan(_UpperCAmelCase			)			)




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



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							Tuple             =   self.model_tester.prepare_config_and_inputs_for_common()
					_UpperCAmelCase		:							List[Any]             =   True
					_UpperCAmelCase		:							List[Any]             =   True

					if hasattr(_UpperCAmelCase						,	"use_cache"			):
							_UpperCAmelCase		:							int             =   True

					_UpperCAmelCase		:							Optional[int]             =   getattr(self.model_tester						,	"encoder_seq_length"						,	self.model_tester.seq_length			)
					_UpperCAmelCase		:							int             =   getattr(self.model_tester						,	"key_length"						,	_UpperCAmelCase			)

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

							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		:							str             =   tf.cast(class_inputs_dict[key]						,	tf.intaa			)

							if set(class_inputs_dict.keys()			) != set(model.input_signature.keys()			):
									continue  # Some models have inputs that the preparation functions don't create, we skip those

							with tempfile.TemporaryDirectory() as tmpdirname:
									model.save_pretrained(_UpperCAmelCase						,	saved_model=_UpperCAmelCase			)
									_UpperCAmelCase		:							Optional[int]             =   os.path.join(_UpperCAmelCase						,	"saved_model"						,	"1"			)
									_UpperCAmelCase		:							Optional[int]             =   tf.keras.models.load_model(_UpperCAmelCase			)
									_UpperCAmelCase		:							Dict             =   model(_UpperCAmelCase			)

									if self.is_encoder_decoder:
											_UpperCAmelCase		:							Any             =   outputs["encoder_hidden_states"]
											_UpperCAmelCase		:							List[str]             =   outputs["encoder_attentions"]
									else:
											_UpperCAmelCase		:							Union[str, Any]             =   outputs["hidden_states"]
											_UpperCAmelCase		:							Any             =   outputs["attentions"]

									self.assertEqual(len(_UpperCAmelCase			)						,	_UpperCAmelCase			)

									_UpperCAmelCase		:							Optional[Any]             =   getattr(
									    self.model_tester						,	"expected_num_hidden_layers"						,	self.model_tester.num_hidden_layers + 1			)

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

									self.assertEqual(len(_UpperCAmelCase			)						,	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 				lowerCAmelCase__     (		self	:  int			) ->str:



					'''simple docstring'''







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

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

									self.assertIsNotNone(_UpperCAmelCase			)
					finally:
							tf.keras.mixed_precision.set_global_policy("float32"			)




			@slow
			def 				lowerCAmelCase__     (		self	:  int			) ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							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"]
					_UpperCAmelCase		:							Dict             =   {
					    module_member
					    for model_class in self.all_model_classes
					    for module in (import_module(model_class.__module__			),)
					    for module_member_name in dir(_UpperCAmelCase			)
					    if module_member_name.endswith("MainLayer"			)
					    for module_member in (getattr(_UpperCAmelCase						,	_UpperCAmelCase			),)
					    if isinstance(_UpperCAmelCase						,	_UpperCAmelCase			)
					    and tf.keras.layers.Layer in module_member.__bases__
					    and getattr(_UpperCAmelCase						,	"_keras_serializable"						,	_UpperCAmelCase			)
					}

					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		:							Union[str, Any]             =   TFSharedEmbeddings(self.model_tester.vocab_size						,	self.model_tester.hidden_size						,	name="shared"			)
									_UpperCAmelCase		:							Optional[Any]             =   False
									_UpperCAmelCase		:							Optional[int]             =   main_layer_class(_UpperCAmelCase						,	embed_tokens=_UpperCAmelCase			)
							else:
									_UpperCAmelCase		:							List[str]             =   main_layer_class(_UpperCAmelCase			)

							_UpperCAmelCase		:							Union[str, Any]             =   {
							    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		:							Tuple             =   self.model_tester.num_labels
							else:
									_UpperCAmelCase		:							Optional[Any]             =   2

							_UpperCAmelCase		:							str             =   tf.data.Dataset.from_tensor_slices(
							    (inputs_dict, np.ones((self.model_tester.batch_size, self.model_tester.seq_length, num_labels, 1)			))			).batch(1			)

							_UpperCAmelCase		:							Any             =   main_layer(_UpperCAmelCase			)[0]
							_UpperCAmelCase		:							List[str]             =   tf.keras.layers.Dense(_UpperCAmelCase						,	activation="softmax"						,	name="outputs"			)(_UpperCAmelCase			)
							_UpperCAmelCase		:							List[str]             =   tf.keras.models.Model(inputs=_UpperCAmelCase						,	outputs=[outputs]			)

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

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




			@slow
			def 				lowerCAmelCase__     (		self	:  Dict			) ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							Union[str, Any]             =   self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes[:2]:
							_UpperCAmelCase		:							Optional[Any]             =   model_class(_UpperCAmelCase			)

							_UpperCAmelCase		:							List[Any]             =   copy.deepcopy(_UpperCAmelCase			)

							if not self.is_encoder_decoder:
									_UpperCAmelCase		:							Optional[Any]             =   inputs["input_ids"]
									del inputs["input_ids"]
							else:
									_UpperCAmelCase		:							Any             =   inputs["input_ids"]
									_UpperCAmelCase		:							Optional[int]             =   inputs.get("decoder_input_ids"						,	_UpperCAmelCase			)
									del inputs["input_ids"]
									inputs.pop("decoder_input_ids"						,	_UpperCAmelCase			)

							if not self.is_encoder_decoder:
									_UpperCAmelCase		:							Dict             =   model.get_input_embeddings()(_UpperCAmelCase			)
							else:
									_UpperCAmelCase		:							List[Any]             =   model.get_input_embeddings()(_UpperCAmelCase			)
									_UpperCAmelCase		:							List[Any]             =   model.get_input_embeddings()(_UpperCAmelCase			)

							_UpperCAmelCase		:							str             =   self._prepare_for_class(_UpperCAmelCase						,	_UpperCAmelCase			)

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

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




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



					'''simple docstring'''







					_UpperCAmelCase		:							Optional[int]             =   []
					if model.config.bos_token_id is not None:
							special_tokens.append(model.config.bos_token_id			)
					if model.config.pad_token_id is not None:
							special_tokens.append(model.config.pad_token_id			)
					if model.config.eos_token_id is not None:
							special_tokens.append(model.config.eos_token_id			)

					# create random bad tokens that are not special tokens
					_UpperCAmelCase		:							int             =   []
					while len(_UpperCAmelCase			) < num_bad_tokens:
							_UpperCAmelCase		:							str             =   tf.squeeze(ids_tensor((1, 1)						,	self.model_tester.vocab_size			)						,	0			).numpy()[0]
							if token not in special_tokens:
									bad_tokens.append(_UpperCAmelCase			)
					return bad_tokens




			def 				lowerCAmelCase__     (		self	:  int						,	lowerCamelCase__	:  List[str]			) ->Tuple:



					'''simple docstring'''







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






			def 				lowerCAmelCase__     (		self	:  Union[str, Any]						,	lowerCamelCase__	:  List[str]						,	lowerCamelCase__	:  Optional[Any]			) ->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(_UpperCAmelCase			)						,	len(_UpperCAmelCase			)			):
											# if tokens match
											if generated_ids_slice[i - len(_UpperCAmelCase			) : i] == bad_word_ids:
													return True
					return False





from ctypes import c_float, sizeof
from enum import Enum
from typing import TYPE_CHECKING, Optional, Union


if TYPE_CHECKING:
							from .. import AutoFeatureExtractor, AutoProcessor, AutoTokenizer  # tests_ignore
class       A__  (		lowerCAmelCase__							):
							lowerCAmelCase__						:		Any            =				c_float





							@property
							def      a__					(					self    :  List[str]							)		->						int:



											"""simple docstring"""


											return sizeof(self.value							)






def 				__SCREAMING_SNAKE_CASE  (						SCREAMING_SNAKE_CASE			: int       ,				SCREAMING_SNAKE_CASE			: int       ,				SCREAMING_SNAKE_CASE			: int = 0     )	->		int:
				# < 0 is possible if using a dynamic axis
				if dimension <= 0:
								__lowercase								=   fixed_dimension

				dimension -= num_token_to_add
				return dimension




def 				__SCREAMING_SNAKE_CASE  (						SCREAMING_SNAKE_CASE			: int       ,				SCREAMING_SNAKE_CASE			: ParameterFormat     )	->		int:
				return num_parameters * dtype.size




def 				__SCREAMING_SNAKE_CASE  (						SCREAMING_SNAKE_CASE			: str     )	->		Optional[Union["AutoTokenizer", "AutoFeatureExtractor", "AutoProcessor"]]:
				from .. import AutoFeatureExtractor, AutoProcessor, AutoTokenizer  # tests_ignore

				try:
								return AutoProcessor.from_pretrained(SCREAMING_SNAKE_CASE     )
				except (ValueError, OSError, KeyError):
								__lowercase						,   __lowercase								=   None, None
								try:
												__lowercase								=   AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE     )
								except (OSError, KeyError):
												pass
								try:
												__lowercase								=   AutoFeatureExtractor.from_pretrained(SCREAMING_SNAKE_CASE     )
								except (OSError, KeyError):
												pass

								if tokenizer is not None and feature_extractor is not None:
												raise ValueError(
												    F"""Couldn't auto-detect preprocessor for {model_name}. Found both a tokenizer and a feature extractor."""     )
								elif tokenizer is None and feature_extractor is None:
												return None
								elif tokenizer is not None:
												return tokenizer
								else:
												return feature_extractor





