"""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:
												UpperCamelCase_       :       Tuple							   =		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_       :       Tuple							   =		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 			_lowercase		:
				lowerCamelCase__						=       None
				lowerCamelCase__						=       ()
				lowerCamelCase__						=       ()
				lowerCamelCase__						=       True
				lowerCamelCase__						=       True
				lowerCamelCase__						=       True
				lowerCamelCase__						=       False

				def        A (     self       :    Tuple       ,    _lowercase       :    int       ,    _lowercase       :    Any       ,    _lowercase       :    str=False     ):


										'''simple docstring'''
										UpperCAmelCase__						=		copy.deepcopy(_lowercase     )

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

										if return_labels:
																if model_class in get_values(_lowercase     ):
																						UpperCAmelCase__						=		tf.ones(self.model_tester.batch_size       ,    dtype=tf.intaa     )
																elif model_class in get_values(_lowercase     ):
																						UpperCAmelCase__						=		tf.zeros(self.model_tester.batch_size       ,    dtype=tf.intaa     )
																						UpperCAmelCase__						=		tf.zeros(self.model_tester.batch_size       ,    dtype=tf.intaa     )
																elif model_class in [
																    *get_values(_lowercase     ),
																    *get_values(_lowercase     ),
																]:
																						UpperCAmelCase__						=		tf.zeros(self.model_tester.batch_size       ,    dtype=tf.intaa     )
																elif model_class in get_values(_lowercase     ):
																						UpperCAmelCase__						=		tf.zeros(self.model_tester.batch_size       ,    dtype=tf.intaa     )
																elif model_class in [
																    *get_values(_lowercase     ),
																    *get_values(_lowercase     ),
																    *get_values(_lowercase     ),
																    *get_values(_lowercase     ),
																    *get_values(_lowercase     ),
																]:
																						UpperCAmelCase__						=		tf.zeros(
																						    (self.model_tester.batch_size, self.model_tester.seq_length)       ,    dtype=tf.intaa     )
										return inputs_dict

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


										'''simple docstring'''
										UpperCAmelCase__						=		self.model_tester.prepare_config_and_inputs_for_common()
										for model_class in self.all_model_classes[:2]:
																UpperCAmelCase__						=		self._prepare_for_class(_lowercase       ,    _lowercase     )
																UpperCAmelCase__						=		model_class(_lowercase     )

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

																UpperCAmelCase__						=		run_in_graph_mode()
																self.assertIsNotNone(_lowercase     )

				@slow
				def        A (     self       :    Dict     ):


										'''simple docstring'''
										UpperCAmelCase__						=		self.model_tester.prepare_config_and_inputs_for_common()
										for model_class in self.all_model_classes[:2]:
																UpperCAmelCase__						=		self._prepare_for_class(_lowercase       ,    _lowercase     )
																UpperCAmelCase__						=		model_class(_lowercase     )

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

																UpperCAmelCase__						=		run_in_graph_mode()
																self.assertIsNotNone(_lowercase     )

				@slow
				def        A (     self       :    Any     ):


										'''simple docstring'''
										UpperCAmelCase__						=		self.model_tester.prepare_config_and_inputs_for_common()
										for model_class in self.all_model_classes[:2]:
																UpperCAmelCase__						=		model_class(_lowercase     )
																if getattr(_lowercase       ,    """hf_compute_loss"""       ,    _lowercase     ):
																						# Test that model correctly compute the loss with kwargs
																						UpperCAmelCase__						=		self._prepare_for_class(inputs_dict.copy()       ,    _lowercase       ,    return_labels=_lowercase     )
																						# Is there a better way to remove these decoder inputs?
																						UpperCAmelCase__						=		{
																						    key: val
																						    for key, val in prepared_for_class.items()
																						    if key not in ("head_mask", "decoder_head_mask", "cross_attn_head_mask", "decoder_input_ids")
																						}

																						UpperCAmelCase__						=		{
																						    "labels",
																						    "label",
																						    "label_ids",
																						    "start_positions",
																						    "start_position",
																						    "end_positions",
																						    "end_position",
																						    "next_sentence_label",
																						}
																						UpperCAmelCase__						=		possible_label_cols.intersection(set(_lowercase     )     )
																						self.assertGreater(len(_lowercase     )       ,    0       ,    msg="""No matching label names found!"""     )
																						UpperCAmelCase__						=		{key: val for key, val in prepared_for_class.items() if key in label_names}
																						UpperCAmelCase__						=		{key: val for key, val in prepared_for_class.items() if key not in label_names}
																						self.assertGreater(len(_lowercase     )       ,    0     )

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

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

																						UpperCAmelCase__						=		history.history["loss"][0]
																						self.assertTrue(not isnan(_lowercase     )     )
																						UpperCAmelCase__						=		history.history["val_loss"][0]
																						self.assertTrue(not isnan(_lowercase     )     )

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


										'''simple docstring'''
										UpperCAmelCase__						=		self.model_tester.prepare_config_and_inputs_for_common()
										UpperCAmelCase__						=		True
										UpperCAmelCase__						=		True

										if hasattr(_lowercase       ,    """use_cache"""     ):
																UpperCAmelCase__						=		True

										UpperCAmelCase__						=		getattr(self.model_tester       ,    """encoder_seq_length"""       ,    self.model_tester.seq_length     )
										UpperCAmelCase__						=		getattr(self.model_tester       ,    """key_length"""       ,    _lowercase     )

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

																for key in list(class_inputs_dict.keys()     ):
																						# Remove keys not in the serving signature, as the SavedModel will not be compiled to deal with them
																						if key not in model.input_signature:
																												del class_inputs_dict[key]
																						# Check it's a tensor, in case the inputs dict has some bools in it too
																						elif isinstance(class_inputs_dict[key]       ,    tf.Tensor     ) and class_inputs_dict[key].dtype.is_integer:
																												UpperCAmelCase__						=		tf.cast(class_inputs_dict[key]       ,    tf.intaa     )

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

																with tempfile.TemporaryDirectory() as tmpdirname:
																						model.save_pretrained(_lowercase       ,    saved_model=_lowercase     )
																						UpperCAmelCase__						=		os.path.join(_lowercase       ,    """saved_model"""       ,    """1"""     )
																						UpperCAmelCase__						=		tf.keras.models.load_model(_lowercase     )
																						UpperCAmelCase__						=		model(_lowercase     )

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

																						self.assertEqual(len(_lowercase     )       ,    _lowercase     )

																						UpperCAmelCase__						=		getattr(
																						    self.model_tester       ,    """expected_num_hidden_layers"""       ,    self.model_tester.num_hidden_layers + 1     )

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

																						self.assertEqual(len(_lowercase     )       ,    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       :    Tuple     ):


										'''simple docstring'''
										tf.keras.mixed_precision.set_global_policy("""mixed_float16"""     )

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

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

				@slow
				def        A (     self       :    str     ):


										'''simple docstring'''
										UpperCAmelCase__						=		self.model_tester.prepare_config_and_inputs_for_common()
										# head_mask and decoder_head_mask has different shapes than other input args
										if "head_mask" in inputs_dict:
																del inputs_dict["head_mask"]
										if "decoder_head_mask" in inputs_dict:
																del inputs_dict["decoder_head_mask"]
										if "cross_attn_head_mask" in inputs_dict:
																del inputs_dict["cross_attn_head_mask"]
										UpperCAmelCase__						=		{
										    module_member
										    for model_class in self.all_model_classes
										    for module in (import_module(model_class.__module__     ),)
										    for module_member_name in dir(_lowercase     )
										    if module_member_name.endswith("""MainLayer"""     )
										    for module_member in (getattr(_lowercase       ,    _lowercase     ),)
										    if isinstance(_lowercase       ,    _lowercase     )
										    and tf.keras.layers.Layer in module_member.__bases__
										    and getattr(_lowercase       ,    """_keras_serializable"""       ,    _lowercase     )
										}

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

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

																if hasattr(self.model_tester       ,    """num_labels"""     ):
																						UpperCAmelCase__						=		self.model_tester.num_labels
																else:
																						UpperCAmelCase__						=		2

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

																UpperCAmelCase__						=		main_layer(_lowercase     )[0]
																UpperCAmelCase__						=		tf.keras.layers.Dense(_lowercase       ,    activation="""softmax"""       ,    name="""outputs"""     )(_lowercase     )
																UpperCAmelCase__						=		tf.keras.models.Model(inputs=_lowercase       ,    outputs=[outputs]     )

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

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

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


										'''simple docstring'''
										UpperCAmelCase__						=		self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes[:2]:
																UpperCAmelCase__						=		model_class(_lowercase     )

																UpperCAmelCase__						=		copy.deepcopy(_lowercase     )

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

																if not self.is_encoder_decoder:
																						UpperCAmelCase__						=		model.get_input_embeddings()(_lowercase     )
																else:
																						UpperCAmelCase__						=		model.get_input_embeddings()(_lowercase     )
																						UpperCAmelCase__						=		model.get_input_embeddings()(_lowercase     )

																UpperCAmelCase__						=		self._prepare_for_class(_lowercase       ,    _lowercase     )

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

																UpperCAmelCase__						=		run_in_graph_mode()
																self.assertIsNotNone(_lowercase     )

				def        A (     self       :    Dict       ,    _lowercase       :    Optional[int]       ,    _lowercase       :    Optional[int]     ):


										'''simple docstring'''
										UpperCAmelCase__						=		[]
										if model.config.bos_token_id is not None:
																special_tokens.append(model.config.bos_token_id     )
										if model.config.pad_token_id is not None:
																special_tokens.append(model.config.pad_token_id     )
										if model.config.eos_token_id is not None:
																special_tokens.append(model.config.eos_token_id     )

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

				def        A (     self       :    int       ,    _lowercase       :    Tuple     ):


										'''simple docstring'''







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

				def        A (     self       :    Optional[Any]       ,    _lowercase       :    Any       ,    _lowercase       :    str     ):


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


"""simple docstring"""







from ..utils import DummyObject, requires_backends


class 			_lowercase		(   metaclass=__lowerCamelCase  ):
				lowerCamelCase__						=       ['torch', 'transformers', 'invisible_watermark']

				def __init__(     self       :    Tuple       ,    *_lowercase       :    int       ,    **_lowercase       :    List[Any]     ):


										'''simple docstring'''
										requires_backends(self       ,    ["""torch""", """transformers""", """invisible_watermark"""]     )

				@classmethod
				def        A (     cls       :    Tuple       ,    *_lowercase       :    Dict       ,    **_lowercase       :    Union[str, Any]     ):


										'''simple docstring'''
										requires_backends(cls       ,    ["""torch""", """transformers""", """invisible_watermark"""]     )

				@classmethod
				def        A (     cls       :    Tuple       ,    *_lowercase       :    Tuple       ,    **_lowercase       :    Any     ):


										'''simple docstring'''
										requires_backends(cls       ,    ["""torch""", """transformers""", """invisible_watermark"""]     )


class 			_lowercase		(   metaclass=__lowerCamelCase  ):
				lowerCamelCase__						=       ['torch', 'transformers', 'invisible_watermark']

				def __init__(     self       :    Union[str, Any]       ,    *_lowercase       :    List[Any]       ,    **_lowercase       :    List[Any]     ):


										'''simple docstring'''
										requires_backends(self       ,    ["""torch""", """transformers""", """invisible_watermark"""]     )

				@classmethod
				def        A (     cls       :    List[str]       ,    *_lowercase       :    Any       ,    **_lowercase       :    Tuple     ):


										'''simple docstring'''
										requires_backends(cls       ,    ["""torch""", """transformers""", """invisible_watermark"""]     )

				@classmethod
				def        A (     cls       :    List[Any]       ,    *_lowercase       :    Optional[Any]       ,    **_lowercase       :    Union[str, Any]     ):


										'''simple docstring'''
										requires_backends(cls       ,    ["""torch""", """transformers""", """invisible_watermark"""]     )

