

"""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:
								_SCREAMING_SNAKE_CASE							:	Union[str, 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)]
																)
																_SCREAMING_SNAKE_CASE							:	List[str]										=			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 :

			"""simple docstring"""

			SCREAMING_SNAKE_CASE_      =    None
			SCREAMING_SNAKE_CASE_      =    ()
			SCREAMING_SNAKE_CASE_      =    ()
			SCREAMING_SNAKE_CASE_      =    True
			SCREAMING_SNAKE_CASE_      =    True
			SCREAMING_SNAKE_CASE_      =    True
			SCREAMING_SNAKE_CASE_      =    False






			def  _lowerCAmelCase      (				self  :       Optional[Any]    ,	lowercase_  :       str    ,	lowercase_  :       List[str]    ,	lowercase_  :       List[str]=False    ):
								UpperCamelCase__		:     Optional[Any]								=copy.deepcopy(lowercase_    )

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






			@slow
			def  _lowerCAmelCase      (				self  :       Optional[Any]    ):
								UpperCamelCase__	,      UpperCamelCase__		:     str								=self.model_tester.prepare_config_and_inputs_for_common()
								for model_class in self.all_model_classes[:2]:
													UpperCamelCase__		:     List[Any]								=self._prepare_for_class(lowercase_    ,	lowercase_    )
													UpperCamelCase__		:     List[Any]								=model_class(lowercase_    )

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

													UpperCamelCase__		:     Optional[int]								=run_in_graph_mode()
													self.assertIsNotNone(lowercase_    )






			@slow
			def  _lowerCAmelCase      (				self  :       int    ):
								UpperCamelCase__	,      UpperCamelCase__		:     List[str]								=self.model_tester.prepare_config_and_inputs_for_common()
								for model_class in self.all_model_classes[:2]:
													UpperCamelCase__		:     int								=self._prepare_for_class(lowercase_    ,	lowercase_    )
													UpperCamelCase__		:     Optional[int]								=model_class(lowercase_    )

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

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






			@slow
			def  _lowerCAmelCase      (				self  :       List[Any]    ):
								# This is a copy of the test_keras_fit method, but we use XLA compilation instead of eager
								UpperCamelCase__	,      UpperCamelCase__		:     Any								=self.model_tester.prepare_config_and_inputs_for_common()
								for model_class in self.all_model_classes[:2]:
													UpperCamelCase__		:     int								=model_class(lowercase_    )
													if getattr(lowercase_    ,	'''hf_compute_loss'''    ,	lowercase_    ):
																		# Test that model correctly compute the loss with kwargs
																		UpperCamelCase__		:     str								=self._prepare_for_class(inputs_dict.copy()    ,	lowercase_    ,	return_labels=lowercase_    )
																		# Is there a better way to remove these decoder inputs?
																		UpperCamelCase__		:     Dict								={
																		    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__		:     Optional[Any]								={
																		    '''labels''',
																		    '''label''',
																		    '''label_ids''',
																		    '''start_positions''',
																		    '''start_position''',
																		    '''end_positions''',
																		    '''end_position''',
																		    '''next_sentence_label''',
																		}
																		UpperCamelCase__		:     List[Any]								=possible_label_cols.intersection(set(lowercase_    )    )
																		self.assertGreater(len(lowercase_    )    ,	0    ,	msg='''No matching label names found!'''    )
																		UpperCamelCase__		:     Union[str, Any]								={key: val for key, val in prepared_for_class.items() if key in label_names}
																		UpperCamelCase__		:     Optional[Any]								={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__		:     List[str]								=model.fit(
																		    lowercase_    ,	validation_data=lowercase_    ,	steps_per_epoch=1    ,	validation_steps=1    ,	shuffle=lowercase_    ,	verbose=0    ,	)
																		UpperCamelCase__		:     Tuple								=history.history['''loss'''][0]
																		self.assertTrue(not isnan(lowercase_    )    )
																		UpperCamelCase__		:     Optional[Any]								=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__		:     Optional[int]								=model_class(lowercase_    )
																		model.compile(optimizer=tf.keras.optimizers.SGD(0.0    )    ,	jit_compile=lowercase_    )
																		UpperCamelCase__		:     List[str]								=model.fit(
																		    lowercase_    ,	lowercase_    ,	validation_data=(inputs_minus_labels, labels)    ,	steps_per_epoch=1    ,	validation_steps=1    ,	shuffle=lowercase_    ,	verbose=0    ,	)

																		UpperCamelCase__		:     Dict								=history.history['''loss'''][0]
																		self.assertTrue(not isnan(lowercase_    )    )
																		UpperCamelCase__		:     Dict								=history.history['''val_loss'''][0]
																		self.assertTrue(not isnan(lowercase_    )    )






			@slow
			def  _lowerCAmelCase      (				self  :       Optional[Any]    ):
								UpperCamelCase__	,      UpperCamelCase__		:     Tuple								=self.model_tester.prepare_config_and_inputs_for_common()
								UpperCamelCase__		:     int								=True
								UpperCamelCase__		:     Union[str, Any]								=True

								if hasattr(lowercase_    ,	'''use_cache'''    ):
													UpperCamelCase__		:     Dict								=True

								UpperCamelCase__		:     List[Any]								=getattr(self.model_tester    ,	'''encoder_seq_length'''    ,	self.model_tester.seq_length    )
								UpperCamelCase__		:     Optional[int]								=getattr(self.model_tester    ,	'''key_length'''    ,	lowercase_    )

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

																		if self.is_encoder_decoder:
																							UpperCamelCase__		:     Optional[int]								=outputs['''encoder_hidden_states''']
																							UpperCamelCase__		:     List[Any]								=outputs['''encoder_attentions''']
																		else:
																							UpperCamelCase__		:     Any								=outputs['''hidden_states''']
																							UpperCamelCase__		:     List[Any]								=outputs['''attentions''']

																		self.assertEqual(len(lowercase_    )    ,	lowercase_    )

																		UpperCamelCase__		:     Tuple								=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  _lowerCAmelCase      (				self  :       Optional[Any]    ):
								tf.keras.mixed_precision.set_global_policy('''mixed_float16'''    )

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

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






			@slow
			def  _lowerCAmelCase      (				self  :       str    ):
								UpperCamelCase__	,      UpperCamelCase__		:     Optional[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__		:     Optional[int]								={
								    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__		:     Dict								=TFSharedEmbeddings(self.model_tester.vocab_size    ,	self.model_tester.hidden_size    ,	name='''shared'''    )
																		UpperCamelCase__		:     str								=False
																		UpperCamelCase__		:     Union[str, Any]								=main_layer_class(lowercase_    ,	embed_tokens=lowercase_    )
													else:
																		UpperCamelCase__		:     int								=main_layer_class(lowercase_    )

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

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

													UpperCamelCase__		:     Dict								=copy.deepcopy(lowercase_    )

													if not self.is_encoder_decoder:
																		UpperCamelCase__		:     Tuple								=inputs['''input_ids''']
																		del inputs["input_ids"]
													else:
																		UpperCamelCase__		:     Union[str, Any]								=inputs['''input_ids''']
																		UpperCamelCase__		:     str								=inputs.get('''decoder_input_ids'''    ,	lowercase_    )
																		del inputs["input_ids"]
																		inputs.pop('''decoder_input_ids'''    ,	lowercase_    )

													if not self.is_encoder_decoder:
																		UpperCamelCase__		:     Dict								=model.get_input_embeddings()(lowercase_    )
													else:
																		UpperCamelCase__		:     int								=model.get_input_embeddings()(lowercase_    )
																		UpperCamelCase__		:     Optional[Any]								=model.get_input_embeddings()(lowercase_    )

													UpperCamelCase__		:     Tuple								=self._prepare_for_class(lowercase_    ,	lowercase_    )

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

													UpperCamelCase__		:     List[str]								=run_in_graph_mode()
													self.assertIsNotNone(lowercase_    )






			def  _lowerCAmelCase      (				self  :       List[str]    ,	lowercase_  :       int    ,	lowercase_  :       Union[str, Any]    ):
								# special tokens cannot be bad tokens
								UpperCamelCase__		:     str								=[]
								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__		:     Optional[Any]								=[]
								while len(lowercase_    ) < num_bad_tokens:
													UpperCamelCase__		:     int								=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  _lowerCAmelCase      (				self  :       Union[str, Any]    ,	lowercase_  :       Dict    ):
								for token_id in output_ids[0].numpy().tolist():
													self.assertGreaterEqual(lowercase_    ,	0    )
													self.assertLess(lowercase_    ,	self.model_tester.vocab_size    )






			def  _lowerCAmelCase      (				self  :       str    ,	lowercase_  :       str    ,	lowercase_  :       Optional[int]    ):
								# for all bad word tokens
								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"""





def 							_lowerCAmelCase   (    UpperCAmelCase  :				int   ,				UpperCAmelCase  :				int   ,				UpperCAmelCase  :				int					):


					'''simple docstring'''






					if power < 0:
										return -1
					base %= mod
					UpperCamelCase__		:     Any								=1

					while power > 0:
										if power & 1:
															UpperCamelCase__		:     Optional[Any]								=(result * base) % mod
										UpperCamelCase__		:     Union[str, Any]								=power >> 1
										UpperCamelCase__		:     Union[str, Any]								=(base * base) % mod

					return result



def 							_lowerCAmelCase   (    ):


					'''simple docstring'''






					print(modular_exponential(3   ,				200   ,				13					)					)


if __name__ == "__main__":
				import doctest

				doctest.testmod()

				main()


