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





	def 							__SCREAMING_SNAKE_CASE					(   self     :    int ,							__magic_name__     :    str ,							__magic_name__     :    Union[str, Any] ,							__magic_name__     :    Optional[Any]=False							):


				"""simple docstring"""

				lowerCAmelCase__									=							copy.deepcopy(UpperCamelCase_							)

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





	@slow
	def 							__SCREAMING_SNAKE_CASE					(   self     :    Optional[Any]							):


				"""simple docstring"""

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

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

							lowerCAmelCase__									=							run_in_graph_mode()
							self.assertIsNotNone(UpperCamelCase_							)





	@slow
	def 							__SCREAMING_SNAKE_CASE					(   self     :    List[Any]							):


				"""simple docstring"""

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

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

							lowerCAmelCase__									=							run_in_graph_mode()
							self.assertIsNotNone(UpperCamelCase_							)





	@slow
	def 							__SCREAMING_SNAKE_CASE					(   self     :    Tuple							):


				"""simple docstring"""

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

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

										lowerCAmelCase__									=							history.history["loss"][0]
										self.assertTrue(not isnan(UpperCamelCase_							)							)
										lowerCAmelCase__									=							history.history["val_loss"][0]
										self.assertTrue(not isnan(UpperCamelCase_							)							)





	@slow
	def 							__SCREAMING_SNAKE_CASE					(   self     :    Optional[Any]							):


				"""simple docstring"""

				lowerCAmelCase__		,lowerCAmelCase__									=							self.model_tester.prepare_config_and_inputs_for_common()
				lowerCAmelCase__									=							True
				lowerCAmelCase__									=							True

				if hasattr(UpperCamelCase_ ,							"use_cache"							):
							lowerCAmelCase__									=							True

				lowerCAmelCase__									=							getattr(self.model_tester ,							"encoder_seq_length" ,							self.model_tester.seq_length							)
				lowerCAmelCase__									=							getattr(self.model_tester ,							"key_length" ,							UpperCamelCase_							)

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

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

										self.assertEqual(len(UpperCamelCase_							) ,							UpperCamelCase_							)

										lowerCAmelCase__									=							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     :    Any							):


				"""simple docstring"""

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

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

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





	@slow
	def 							__SCREAMING_SNAKE_CASE					(   self     :    Any							):


				"""simple docstring"""

				lowerCAmelCase__		,lowerCAmelCase__									=							self.model_tester.prepare_config_and_inputs_for_common()
				# head_mask and decoder_head_mask has different shapes than other input args
				if "head_mask" in inputs_dict:
							del inputs_dict["head_mask"]
				if "decoder_head_mask" in inputs_dict:
							del inputs_dict["decoder_head_mask"]
				if "cross_attn_head_mask" in inputs_dict:
							del inputs_dict["cross_attn_head_mask"]
				lowerCAmelCase__									=							{
				    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
										lowerCAmelCase__									=							TFSharedEmbeddings(self.model_tester.vocab_size ,							self.model_tester.hidden_size ,							name="shared"							)
										lowerCAmelCase__									=							False
										lowerCAmelCase__									=							main_layer_class(UpperCamelCase_ ,							embed_tokens=UpperCamelCase_							)
							else:
										lowerCAmelCase__									=							main_layer_class(UpperCamelCase_							)

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

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

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

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


				"""simple docstring"""

				lowerCAmelCase__		,lowerCAmelCase__									=							self.model_tester.prepare_config_and_inputs_for_common()

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

							lowerCAmelCase__									=							copy.deepcopy(UpperCamelCase_							)

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

							if not self.is_encoder_decoder:
										lowerCAmelCase__									=							model.get_input_embeddings()(UpperCamelCase_							)
							else:
										lowerCAmelCase__									=							model.get_input_embeddings()(UpperCamelCase_							)
										lowerCAmelCase__									=							model.get_input_embeddings()(UpperCamelCase_							)

							lowerCAmelCase__									=							self._prepare_for_class(UpperCamelCase_ ,							UpperCamelCase_							)

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

							lowerCAmelCase__									=							run_in_graph_mode()
							self.assertIsNotNone(UpperCamelCase_							)





	def 							__SCREAMING_SNAKE_CASE					(   self     :    Tuple ,							__magic_name__     :    Dict ,							__magic_name__     :    str							):


				"""simple docstring"""

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

				# create random bad tokens that are not special tokens
				lowerCAmelCase__									=							[]
				while len(UpperCamelCase_							) < num_bad_tokens:
							lowerCAmelCase__									=							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     :    Dict ,							__magic_name__     :    int							):


				"""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     :    List[Any] ,							__magic_name__     :    Optional[int] ,							__magic_name__     :    Union[str, 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







'''simple docstring'''

# this script builds a small sample spm file tests/fixtures/test_sentencepiece_no_bos.model, with features needed by pegasus

# 1. pip install sentencepiece
#
# 2. wget https://raw.githubusercontent.com/google/sentencepiece/master/data/botchan.txt

# 3. build
import sentencepiece as spm

# pegasus:
# 1. no bos
# 2. eos_id is 1
# 3. unk_id is 2
# build a sample spm file accordingly
spm.SentencePieceTrainer.train(
    "--input=botchan.txt --model_prefix=test_sentencepiece_no_bos --bos_id=-1 --unk_id=2  --eos_id=1  --vocab_size=1000"
)

# 4. now update the fixture
# mv test_sentencepiece_no_bos.model ../../tests/fixtures/





