



"""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:
						A    =	tf.config.list_physical_devices('GPU')
						for gpu in gpus:
									# Restrict TensorFlow to only allocate x GB of memory on the GPUs
									try:
												tf.config.set_logical_device_configuration(
												    gpu, [tf.config.LogicalDeviceConfiguration(memory_limit=_tf_gpu_memory_limit)]
												)
												A    =	tf.config.list_logical_devices('GPU')
												print('Logical GPUs', logical_gpus)
									except RuntimeError as e:
												# Virtual devices must be set before GPUs have been initialized
												print(e)

@require_tf
class 			_a     :
			__magic_name__    				=    None
			__magic_name__    				=    ()
			__magic_name__    				=    ()
			__magic_name__    				=    True
			__magic_name__    				=    True
			__magic_name__    				=    True
			__magic_name__    				=    False



			def 							__lowercase				( self       :							Union[str, Any]      ,  _lowercase       :							int      ,  _lowercase       :							Dict      ,  _lowercase       :							Dict=False		)       ->				dict:
									snake_case      : List[str]							=   copy.deepcopy(_lowercase		)

									if model_class in get_values(_lowercase		):
															snake_case      : 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		):
																					snake_case      : Any							=   tf.ones(self.model_tester.batch_size      ,  dtype=tf.intaa		)
															elif model_class in get_values(_lowercase		):
																					snake_case      : Optional[Any]							=   tf.zeros(self.model_tester.batch_size      ,  dtype=tf.intaa		)
																					snake_case      : List[Any]							=   tf.zeros(self.model_tester.batch_size      ,  dtype=tf.intaa		)
															elif model_class in [
															    *get_values(_lowercase		),
															    *get_values(_lowercase		),
															]:
																					snake_case      : Optional[int]							=   tf.zeros(self.model_tester.batch_size      ,  dtype=tf.intaa		)
															elif model_class in get_values(_lowercase		):
																					snake_case      : Optional[Any]							=   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		),
															]:
																					snake_case      : int							=   tf.zeros(
																					    (self.model_tester.batch_size, self.model_tester.seq_length)      ,  dtype=tf.intaa		)
									return inputs_dict



			@slow
			def 							__lowercase				( self       :							str		)       ->				int:
									snake_case						, snake_case      : Dict							=   self.model_tester.prepare_config_and_inputs_for_common()
									for model_class in self.all_model_classes[:2]:
															snake_case      : Optional[int]							=   self._prepare_for_class(_lowercase      ,  _lowercase		)
															snake_case      : Optional[int]							=   model_class(_lowercase		)

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

															snake_case      : Any							=   run_in_graph_mode()
															self.assertIsNotNone(_lowercase		)



			@slow
			def 							__lowercase				( self       :							List[Any]		)       ->				List[str]:
									snake_case						, snake_case      : List[Any]							=   self.model_tester.prepare_config_and_inputs_for_common()
									for model_class in self.all_model_classes[:2]:
															snake_case      : Optional[int]							=   self._prepare_for_class(_lowercase      ,  _lowercase		)
															snake_case      : str							=   model_class(_lowercase		)

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

															snake_case      : str							=   run_in_graph_mode()
															self.assertIsNotNone(_lowercase		)



			@slow
			def 							__lowercase				( self       :							Union[str, Any]		)       ->				Dict:
									# This is a copy of the test_keras_fit method, but we use XLA compilation instead of eager
									snake_case						, snake_case      : Optional[int]							=   self.model_tester.prepare_config_and_inputs_for_common()
									for model_class in self.all_model_classes[:2]:
															snake_case      : Any							=   model_class(_lowercase		)
															if getattr(_lowercase      ,  "hf_compute_loss"      ,  _lowercase		):
																					# Test that model correctly compute the loss with kwargs
																					snake_case      : Optional[Any]							=   self._prepare_for_class(inputs_dict.copy()      ,  _lowercase      ,  return_labels=_lowercase		)
																					# Is there a better way to remove these decoder inputs?
																					snake_case      : List[str]							=   {
																					    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")
																					}

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

																					snake_case      : List[str]							=   history.history["loss"][0]
																					self.assertTrue(not isnan(_lowercase		)		)
																					snake_case      : Optional[Any]							=   history.history["val_loss"][0]
																					self.assertTrue(not isnan(_lowercase		)		)



			@slow
			def 							__lowercase				( self       :							Optional[int]		)       ->				str:
									snake_case						, snake_case      : Any							=   self.model_tester.prepare_config_and_inputs_for_common()
									snake_case      : Union[str, Any]							=   True
									snake_case      : List[str]							=   True

									if hasattr(_lowercase      ,  "use_cache"		):
															snake_case      : List[str]							=   True

									snake_case      : List[str]							=   getattr(self.model_tester      ,  "encoder_seq_length"      ,  self.model_tester.seq_length		)
									snake_case      : int							=   getattr(self.model_tester      ,  "key_length"      ,  _lowercase		)

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

																					if self.is_encoder_decoder:
																											snake_case      : Union[str, Any]							=   outputs["encoder_hidden_states"]
																											snake_case      : Optional[int]							=   outputs["encoder_attentions"]
																					else:
																											snake_case      : str							=   outputs["hidden_states"]
																											snake_case      : Any							=   outputs["attentions"]

																					self.assertEqual(len(_lowercase		)      ,  _lowercase		)

																					snake_case      : int							=   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 							__lowercase				( self       :							str		)       ->				Optional[int]:
									tf.keras.mixed_precision.set_global_policy("mixed_float16"		)

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

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



			@slow
			def 							__lowercase				( self       :							Optional[int]		)       ->				Union[str, Any]:
									snake_case						, snake_case      : Optional[int]							=   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"]
									snake_case      : 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
																					snake_case      : Tuple							=   TFSharedEmbeddings(self.model_tester.vocab_size      ,  self.model_tester.hidden_size      ,  name="shared"		)
																					snake_case      : Any							=   False
																					snake_case      : Tuple							=   main_layer_class(_lowercase      ,  embed_tokens=_lowercase		)
															else:
																					snake_case      : int							=   main_layer_class(_lowercase		)

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

															if hasattr(self.model_tester      ,  "num_labels"		):
																					snake_case      : Union[str, Any]							=   self.model_tester.num_labels
															else:
																					snake_case      : int							=   2

															snake_case      : Union[str, Any]							=   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		)

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

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

															snake_case      : Union[str, Any]							=   copy.deepcopy(_lowercase		)

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

															if not self.is_encoder_decoder:
																					snake_case      : Optional[int]							=   model.get_input_embeddings()(_lowercase		)
															else:
																					snake_case      : Optional[int]							=   model.get_input_embeddings()(_lowercase		)
																					snake_case      : Tuple							=   model.get_input_embeddings()(_lowercase		)

															snake_case      : Any							=   self._prepare_for_class(_lowercase      ,  _lowercase		)

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

															snake_case      : Dict							=   run_in_graph_mode()
															self.assertIsNotNone(_lowercase		)



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



			def 							__lowercase				( self       :							str      ,  _lowercase       :							Dict      ,  _lowercase       :							str		)       ->				Tuple:
									# 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"""







from ...configuration_utils import PretrainedConfig
from ...utils import logging


A    =	logging.get_logger(__name__)

A    =	{
    'uw-madison/yoso-4096': 'https://huggingface.co/uw-madison/yoso-4096/resolve/main/config.json',
    # See all YOSO models at https://huggingface.co/models?filter=yoso
}

class 			_a     (   SCREAMING_SNAKE_CASE__):
			__magic_name__    				=    """yoso"""



			def __init__( self       :							List[str]      ,  _lowercase       :							List[str]=50265      ,  _lowercase       :							List[str]=768      ,  _lowercase       :							Optional[int]=12      ,  _lowercase       :							str=12      ,  _lowercase       :							List[Any]=3072      ,  _lowercase       :							Dict="gelu"      ,  _lowercase       :							int=0.1      ,  _lowercase       :							Optional[Any]=0.1      ,  _lowercase       :							Dict=4096      ,  _lowercase       :							Optional[Any]=1      ,  _lowercase       :							Any=0.02      ,  _lowercase       :							Union[str, Any]=1E-12      ,  _lowercase       :							int="absolute"      ,  _lowercase       :							Any=True      ,  _lowercase       :							Dict=9      ,  _lowercase       :							int=64      ,  _lowercase       :							Any=None      ,  _lowercase       :							List[str]=True      ,  _lowercase       :							Optional[Any]=True      ,  _lowercase       :							Any=1      ,  _lowercase       :							int=0      ,  _lowercase       :							List[str]=2      ,  **_lowercase       :							List[Any]      ,  )       ->				Union[str, Any]:
									super().__init__(pad_token_id=_lowercase      ,  bos_token_id=_lowercase      ,  eos_token_id=_lowercase      ,  **_lowercase		)

									snake_case      : List[str]							=   vocab_size
									snake_case      : Tuple							=   max_position_embeddings
									snake_case      : Optional[Any]							=   hidden_size
									snake_case      : Tuple							=   num_hidden_layers
									snake_case      : Tuple							=   num_attention_heads
									snake_case      : Any							=   intermediate_size
									snake_case      : Union[str, Any]							=   hidden_act
									snake_case      : int							=   hidden_dropout_prob
									snake_case      : Dict							=   attention_probs_dropout_prob
									snake_case      : Any							=   initializer_range
									snake_case      : Optional[Any]							=   type_vocab_size
									snake_case      : List[str]							=   layer_norm_eps
									snake_case      : Dict							=   position_embedding_type
									snake_case      : Tuple							=   use_expectation
									snake_case      : Dict							=   hash_code_len
									snake_case      : Optional[Any]							=   num_hash
									snake_case      : List[str]							=   conv_window
									snake_case      : Optional[Any]							=   use_fast_hash
									snake_case      : Tuple							=   lsh_backward



