

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:
												A__							=				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)]
																								)
																								A__							=				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 					__UpperCamelCase					:
	_lowercase	: List[str]												=	None
	_lowercase	: str												=	()
	_lowercase	: str												=	()
	_lowercase	: Optional[Any]												=	True
	_lowercase	: Union[str, Any]												=	True
	_lowercase	: List[str]												=	True
	_lowercase	: Any												=	False







	def 						_SCREAMING_SNAKE_CASE				( self:			int	,			__UpperCamelCase:			Tuple	,			__UpperCamelCase:			Union[str, Any]	,			__UpperCamelCase:			Dict=False   ):

						'''simple docstring'''





						__magic_name__				=     copy.deepcopy(UpperCamelCase__   )

						if model_class in get_values(UpperCamelCase__   ):
											__magic_name__				=     {
											    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__   ):
																__magic_name__				=     tf.ones(self.model_tester.batch_size	,			dtype=tf.intaa   )
											elif model_class in get_values(UpperCamelCase__   ):
																__magic_name__				=     tf.zeros(self.model_tester.batch_size	,			dtype=tf.intaa   )
																__magic_name__				=     tf.zeros(self.model_tester.batch_size	,			dtype=tf.intaa   )
											elif model_class in [
											    *get_values(UpperCamelCase__   ),
											    *get_values(UpperCamelCase__   ),
											]:
																__magic_name__				=     tf.zeros(self.model_tester.batch_size	,			dtype=tf.intaa   )
											elif model_class in get_values(UpperCamelCase__   ):
																__magic_name__				=     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__   ),
											]:
																__magic_name__				=     tf.zeros(
																    (self.model_tester.batch_size, self.model_tester.seq_length)	,			dtype=tf.intaa   )
						return inputs_dict







	@slow
	def 						_SCREAMING_SNAKE_CASE				( self:			int   ):

						'''simple docstring'''





						__magic_name__				=     self.model_tester.prepare_config_and_inputs_for_common()
						for model_class in self.all_model_classes[:2]:
											__magic_name__				=     self._prepare_for_class(UpperCamelCase__	,			UpperCamelCase__   )
											__magic_name__				=     model_class(UpperCamelCase__   )

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

											__magic_name__				=     run_in_graph_mode()
											self.assertIsNotNone(UpperCamelCase__   )







	@slow
	def 						_SCREAMING_SNAKE_CASE				( self:			Dict   ):

						'''simple docstring'''





						__magic_name__				=     self.model_tester.prepare_config_and_inputs_for_common()
						for model_class in self.all_model_classes[:2]:
											__magic_name__				=     self._prepare_for_class(UpperCamelCase__	,			UpperCamelCase__   )
											__magic_name__				=     model_class(UpperCamelCase__   )

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

											__magic_name__				=     run_in_graph_mode()
											self.assertIsNotNone(UpperCamelCase__   )







	@slow
	def 						_SCREAMING_SNAKE_CASE				( self:			int   ):

						'''simple docstring'''





						__magic_name__				=     self.model_tester.prepare_config_and_inputs_for_common()
						for model_class in self.all_model_classes[:2]:
											__magic_name__				=     model_class(UpperCamelCase__   )
											if getattr(UpperCamelCase__	,			'hf_compute_loss'	,			UpperCamelCase__   ):
																# Test that model correctly compute the loss with kwargs
																__magic_name__				=     self._prepare_for_class(inputs_dict.copy()	,			UpperCamelCase__	,			return_labels=UpperCamelCase__   )
																# Is there a better way to remove these decoder inputs?
																__magic_name__				=     {
																    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""")
																}

																__magic_name__				=     {
																    """labels""",
																    """label""",
																    """label_ids""",
																    """start_positions""",
																    """start_position""",
																    """end_positions""",
																    """end_position""",
																    """next_sentence_label""",
																}
																__magic_name__				=     possible_label_cols.intersection(set(UpperCamelCase__   )   )
																self.assertGreater(len(UpperCamelCase__   )	,			0	,			msg='No matching label names found!'   )
																__magic_name__				=     {key: val for key, val in prepared_for_class.items() if key in label_names}
																__magic_name__				=     {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
																__magic_name__				=     model.fit(
																    UpperCamelCase__	,			validation_data=UpperCamelCase__	,			steps_per_epoch=1	,			validation_steps=1	,			shuffle=UpperCamelCase__	,			verbose=0	,			)
																__magic_name__				=     history.history["""loss"""][0]
																self.assertTrue(not isnan(UpperCamelCase__   )   )
																__magic_name__				=     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.
																__magic_name__				=     model_class(UpperCamelCase__   )
																model.compile(optimizer=tf.keras.optimizers.SGD(0.0   )	,			jit_compile=UpperCamelCase__   )
																__magic_name__				=     model.fit(
																    UpperCamelCase__	,			UpperCamelCase__	,			validation_data=(inputs_minus_labels, labels)	,			steps_per_epoch=1	,			validation_steps=1	,			shuffle=UpperCamelCase__	,			verbose=0	,			)

																__magic_name__				=     history.history["""loss"""][0]
																self.assertTrue(not isnan(UpperCamelCase__   )   )
																__magic_name__				=     history.history["""val_loss"""][0]
																self.assertTrue(not isnan(UpperCamelCase__   )   )







	@slow
	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[int]   ):

						'''simple docstring'''





						__magic_name__				=     self.model_tester.prepare_config_and_inputs_for_common()
						__magic_name__				=     True
						__magic_name__				=     True

						if hasattr(UpperCamelCase__	,			'use_cache'   ):
											__magic_name__				=     True

						__magic_name__				=     getattr(self.model_tester	,			'encoder_seq_length'	,			self.model_tester.seq_length   )
						__magic_name__				=     getattr(self.model_tester	,			'key_length'	,			UpperCamelCase__   )

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

																if self.is_encoder_decoder:
																					__magic_name__				=     outputs["""encoder_hidden_states"""]
																					__magic_name__				=     outputs["""encoder_attentions"""]
																else:
																					__magic_name__				=     outputs["""hidden_states"""]
																					__magic_name__				=     outputs["""attentions"""]

																self.assertEqual(len(UpperCamelCase__   )	,			UpperCamelCase__   )

																__magic_name__				=     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 						_SCREAMING_SNAKE_CASE				( self:			Dict   ):

						'''simple docstring'''





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

						# try/finally block to ensure subsequent tests run in float32
						try:
											__magic_name__				=     self.model_tester.prepare_config_and_inputs_for_common()
											for model_class in self.all_model_classes[:2]:
																__magic_name__				=     self._prepare_for_class(UpperCamelCase__	,			UpperCamelCase__   )
																__magic_name__				=     model_class(UpperCamelCase__   )
																__magic_name__				=     model(UpperCamelCase__   )

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







	@slow
	def 						_SCREAMING_SNAKE_CASE				( self:			List[str]   ):

						'''simple docstring'''





						__magic_name__				=     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"]
						__magic_name__				=     {
						    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
																__magic_name__				=     TFSharedEmbeddings(self.model_tester.vocab_size	,			self.model_tester.hidden_size	,			name='shared'   )
																__magic_name__				=     False
																__magic_name__				=     main_layer_class(UpperCamelCase__	,			embed_tokens=UpperCamelCase__   )
											else:
																__magic_name__				=     main_layer_class(UpperCamelCase__   )

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

											if hasattr(self.model_tester	,			'num_labels'   ):
																__magic_name__				=     self.model_tester.num_labels
											else:
																__magic_name__				=     2

											__magic_name__				=     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   )

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

						'''simple docstring'''





						__magic_name__				=     self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes[:2]:
											__magic_name__				=     model_class(UpperCamelCase__   )

											__magic_name__				=     copy.deepcopy(UpperCamelCase__   )

											if not self.is_encoder_decoder:
																__magic_name__				=     inputs["""input_ids"""]
																del inputs["input_ids"]
											else:
																__magic_name__				=     inputs["""input_ids"""]
																__magic_name__				=     inputs.get('decoder_input_ids'	,			UpperCamelCase__   )
																del inputs["input_ids"]
																inputs.pop('decoder_input_ids'	,			UpperCamelCase__   )

											if not self.is_encoder_decoder:
																__magic_name__				=     model.get_input_embeddings()(UpperCamelCase__   )
											else:
																__magic_name__				=     model.get_input_embeddings()(UpperCamelCase__   )
																__magic_name__				=     model.get_input_embeddings()(UpperCamelCase__   )

											__magic_name__				=     self._prepare_for_class(UpperCamelCase__	,			UpperCamelCase__   )

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

											__magic_name__				=     run_in_graph_mode()
											self.assertIsNotNone(UpperCamelCase__   )







	def 						_SCREAMING_SNAKE_CASE				( self:			Tuple	,			__UpperCamelCase:			Any	,			__UpperCamelCase:			Optional[Any]   ):

						'''simple docstring'''





						__magic_name__				=     []
						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
						__magic_name__				=     []
						while len(UpperCamelCase__   ) < num_bad_tokens:
											__magic_name__				=     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 						_SCREAMING_SNAKE_CASE				( self:			Union[str, Any]	,			__UpperCamelCase:			List[Any]   ):

						'''simple docstring'''





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







	def 						_SCREAMING_SNAKE_CASE				( self:			Dict	,			__UpperCamelCase:			List[Any]	,			__UpperCamelCase:			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(UpperCamelCase__   )	,			len(UpperCamelCase__   )   ):
																					# if tokens match
																					if generated_ids_slice[i - len(UpperCamelCase__   ) : i] == bad_word_ids:
																										return True
						return False












'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_A      :     Optional[int]       	={'''configuration_visual_bert''': ['''VISUAL_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''VisualBertConfig''']}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_A      :     Optional[int]       	=[
			    '''VISUAL_BERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
			    '''VisualBertForMultipleChoice''',
			    '''VisualBertForPreTraining''',
			    '''VisualBertForQuestionAnswering''',
			    '''VisualBertForRegionToPhraseAlignment''',
			    '''VisualBertForVisualReasoning''',
			    '''VisualBertLayer''',
			    '''VisualBertModel''',
			    '''VisualBertPreTrainedModel''',
			]


if TYPE_CHECKING:
			from .configuration_visual_bert import VISUAL_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, VisualBertConfig

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_visual_bert import (
						    VISUAL_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
						    VisualBertForMultipleChoice,
						    VisualBertForPreTraining,
						    VisualBertForQuestionAnswering,
						    VisualBertForRegionToPhraseAlignment,
						    VisualBertForVisualReasoning,
						    VisualBertLayer,
						    VisualBertModel,
						    VisualBertPreTrainedModel,
						)


else:
			import sys

			_A      :     int       	=_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)

