
'''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 :  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)]
																								)
																								lowerCAmelCase :  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        SCREAMING_SNAKE_CASE__	:
				lowerCAmelCase_			       =				None
				lowerCAmelCase_			       =				()
				lowerCAmelCase_			       =				()
				lowerCAmelCase_			       =				True
				lowerCAmelCase_			       =				True
				lowerCAmelCase_			       =				True
				lowerCAmelCase_			       =				False







				def 		UpperCAmelCase_  (    self		,							A_		,							A_		,							A_=False			)->						dict:







											'''simple docstring'''






											UpperCamelCase					       =			copy.deepcopy(A_			)

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

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







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







											'''simple docstring'''






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

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

																		UpperCamelCase					       =			run_in_graph_mode()
																		self.assertIsNotNone(A_			)







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







											'''simple docstring'''






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

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

																		UpperCamelCase					       =			run_in_graph_mode()
																		self.assertIsNotNone(A_			)







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







											'''simple docstring'''






											UpperCamelCase       ,     UpperCamelCase					       =			self.model_tester.prepare_config_and_inputs_for_common()
											for model_class in self.all_model_classes[:2]:
																		UpperCamelCase					       =			model_class(A_			)
																		if getattr(A_		,							'hf_compute_loss'		,							A_			):
																									# Test that model correctly compute the loss with kwargs
																									UpperCamelCase					       =			self._prepare_for_class(inputs_dict.copy()		,							A_		,							return_labels=A_			)
																									# 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(A_			)			)
																									self.assertGreater(len(A_			)		,							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(A_			)		,							0			)

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

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

																									UpperCamelCase					       =			history.history['loss'][0]
																									self.assertTrue(not isnan(A_			)			)
																									UpperCamelCase					       =			history.history['val_loss'][0]
																									self.assertTrue(not isnan(A_			)			)







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







											'''simple docstring'''






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

											if hasattr(A_		,							'use_cache'			):
																		UpperCamelCase					       =			True

											UpperCamelCase					       =			getattr(self.model_tester		,							'encoder_seq_length'		,							self.model_tester.seq_length			)
											UpperCamelCase					       =			getattr(self.model_tester		,							'key_length'		,							A_			)

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

																		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(A_		,							saved_model=A_			)
																									UpperCamelCase					       =			os.path.join(A_		,							'saved_model'		,							'1'			)
																									UpperCamelCase					       =			tf.keras.models.load_model(A_			)
																									UpperCamelCase					       =			model(A_			)

																									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(A_			)		,							A_			)

																									UpperCamelCase					       =			getattr(
																									    self.model_tester		,							'expected_num_hidden_layers'		,							self.model_tester.num_hidden_layers + 1			)

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

																									self.assertEqual(len(A_			)		,							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 		UpperCAmelCase_  (    self			)->						Optional[int]:







											'''simple docstring'''






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

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

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







				@slow
				def 		UpperCAmelCase_  (    self			)->						Tuple:







											'''simple docstring'''






											UpperCamelCase       ,     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(A_			)
											    if module_member_name.endswith('MainLayer'			)
											    for module_member in (getattr(A_		,							A_			),)
											    if isinstance(A_		,							A_			)
											    and tf.keras.layers.Layer in module_member.__bases__
											    and getattr(A_		,							'_keras_serializable'		,							A_			)
											}

											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(A_		,							embed_tokens=A_			)
																		else:
																									UpperCamelCase					       =			main_layer_class(A_			)

																		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(A_			)[0]
																		UpperCamelCase					       =			tf.keras.layers.Dense(A_		,							activation='softmax'		,							name='outputs'			)(A_			)
																		UpperCamelCase					       =			tf.keras.models.Model(inputs=A_		,							outputs=[outputs]			)

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

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







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







											'''simple docstring'''






											UpperCamelCase       ,     UpperCamelCase					       =			self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes[:2]:
																		UpperCamelCase					       =			model_class(A_			)

																		UpperCamelCase					       =			copy.deepcopy(A_			)

																		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'		,							A_			)
																									del inputs["input_ids"]
																									inputs.pop('decoder_input_ids'		,							A_			)

																		if not self.is_encoder_decoder:
																									UpperCamelCase					       =			model.get_input_embeddings()(A_			)
																		else:
																									UpperCamelCase					       =			model.get_input_embeddings()(A_			)
																									UpperCamelCase					       =			model.get_input_embeddings()(A_			)

																		UpperCamelCase					       =			self._prepare_for_class(A_		,							A_			)

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

																		UpperCamelCase					       =			run_in_graph_mode()
																		self.assertIsNotNone(A_			)







				def 		UpperCAmelCase_  (    self		,							A_		,							A_			)->						Optional[Any]:







											'''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(A_			) < 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(A_			)
											return bad_tokens







				def 		UpperCAmelCase_  (    self		,							A_			)->						Optional[int]:







											'''simple docstring'''






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







				def 		UpperCAmelCase_  (    self		,							A_		,							A_			)->						Dict:







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






'''simple docstring'''
import random
from typing import List, Optional, Union

import numpy as np
import torch

from ..state import AcceleratorState
from .constants import CUDA_DISTRIBUTED_TYPES
from .dataclasses import DistributedType, RNGType
from .imports import is_npu_available, is_tpu_available, is_xpu_available


if is_tpu_available(check_device=False):
						import torch_xla.core.xla_model as xm







def 						A_(			A		:						int						,			A		:						bool = False):
							if device_specific:
														seed += AcceleratorState().process_index
							random.seed(A)
							np.random.seed(A)
							torch.manual_seed(A)
							if is_xpu_available():
														torch.xpu.manual_seed_all(A)
							elif is_npu_available():
														torch.npu.manual_seed_all(A)
							else:
														torch.cuda.manual_seed_all(A)
							# ^^ safe to call this function even if cuda is not available
							if is_tpu_available():
														xm.set_rng_state(A)







def 						A_(			A		:						Optional[RNGType] = None						,			A		:						Optional[torch.Generator] = None):
							# Get the proper rng state
							if rng_type == RNGType.TORCH:
														UpperCamelCase					       =			torch.get_rng_state()
							elif rng_type == RNGType.CUDA:
														UpperCamelCase					       =			torch.cuda.get_rng_state()
							elif rng_type == RNGType.XLA:
														assert is_tpu_available(), "Can't synchronize XLA seeds on an environment without TPUs."
														UpperCamelCase					       =			torch.tensor(xm.get_rng_state())
							elif rng_type == RNGType.NPU:
														assert is_npu_available(), "Can't synchronize NPU seeds on an environment without NPUs."
														UpperCamelCase					       =			torch.npu.get_rng_state()
							elif rng_type == RNGType.XPU:
														assert is_xpu_available(), "Can't synchronize XPU seeds on an environment without XPUs."
														UpperCamelCase					       =			torch.xpu.get_rng_state()
							elif rng_type == RNGType.GENERATOR:
														assert generator is not None, "Need a generator to synchronize its seed."
														UpperCamelCase					       =			generator.get_state()

							# Broadcast the rng state from device 0 to other devices
							UpperCamelCase					       =			AcceleratorState()
							if state.distributed_type == DistributedType.TPU:
														UpperCamelCase					       =			rng_state.to(xm.xla_device())
														xm.collective_broadcast([rng_state])
														xm.mark_step()
														UpperCamelCase					       =			rng_state.cpu()
							elif (
							    state.distributed_type in CUDA_DISTRIBUTED_TYPES
							    or state.distributed_type == DistributedType.MULTI_NPU
							    or state.distributed_type == DistributedType.MULTI_XPU
							):
														UpperCamelCase					       =			rng_state.to(state.device)
														torch.distributed.broadcast(A						,			0)
														UpperCamelCase					       =			rng_state.cpu()
							elif state.distributed_type == DistributedType.MULTI_CPU:
														torch.distributed.broadcast(A						,			0)

							# Set the broadcast rng state
							if rng_type == RNGType.TORCH:
														torch.set_rng_state(A)
							elif rng_type == RNGType.CUDA:
														torch.cuda.set_rng_state(A)
							elif rng_type == RNGType.NPU:
														torch.npu.set_rng_state(A)
							elif rng_type == RNGType.XPU:
														torch.xpu.set_rng_state(A)
							elif rng_type == RNGType.XLA:
														xm.set_rng_state(rng_state.item())
							elif rng_type == RNGType.GENERATOR:
														generator.set_state(A)







def 						A_(			A		:						List[Union[str, RNGType]]						,			A		:						Optional[torch.Generator] = None):
							for rng_type in rng_types:
														synchronize_rng_state(RNGType(A)						,			generator=A)




