
"""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 	lowerCamelCase__      :





			'''simple docstring'''
			lowerCamelCase						      =	None
			lowerCamelCase						      =	()
			lowerCamelCase						      =	()
			lowerCamelCase						      =	True
			lowerCamelCase						      =	True
			lowerCamelCase						      =	True
			lowerCamelCase						      =	False



			def       _lowerCAmelCase	(							self  ,  __UpperCAmelCase  ,  __UpperCAmelCase  ,  __UpperCAmelCase=False						)				->  dict:
				_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       _lowerCAmelCase	(							self						)				->  Tuple:
				_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       _lowerCAmelCase	(							self						)				->  List[str]:
				_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       _lowerCAmelCase	(							self						)				->  Union[str, Any]:
				# This is a copy of the test_keras_fit method, but we use XLA compilation instead of eager
				_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       _lowerCAmelCase	(							self						)				->  int:
				_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       _lowerCAmelCase	(							self						)				->  List[str]:
				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       _lowerCAmelCase	(							self						)				->  List[str]:
				_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       _lowerCAmelCase	(							self						)				->  Tuple:
				_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       _lowerCAmelCase	(							self  ,  __UpperCAmelCase  ,  __UpperCAmelCase						)				->  Union[str, Any]:
				# special tokens cannot be bad tokens
				_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       _lowerCAmelCase	(							self  ,  __UpperCAmelCase						)				->  Optional[int]:
				for token_id in output_ids[0].numpy().tolist():
					self.assertGreaterEqual(__UpperCAmelCase  ,  0						)
					self.assertLess(__UpperCAmelCase  ,  self.model_tester.vocab_size						)



			def       _lowerCAmelCase	(							self  ,  __UpperCAmelCase  ,  __UpperCAmelCase						)				->  List[str]:
				# 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(__UpperCAmelCase						)  ,  len(__UpperCAmelCase						)						):
							# if tokens match
							if generated_ids_slice[i - len(__UpperCAmelCase						) : i] == bad_word_ids:
								return True
				return False

"""simple docstring"""


from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__A 		= logging.get_logger(__name__)

__A 		= {
    'nghuyong/ernie-1.0-base-zh': 'https://huggingface.co/nghuyong/ernie-1.0-base-zh/resolve/main/config.json',
    'nghuyong/ernie-2.0-base-en': 'https://huggingface.co/nghuyong/ernie-2.0-base-en/resolve/main/config.json',
    'nghuyong/ernie-2.0-large-en': 'https://huggingface.co/nghuyong/ernie-2.0-large-en/resolve/main/config.json',
    'nghuyong/ernie-3.0-base-zh': 'https://huggingface.co/nghuyong/ernie-3.0-base-zh/resolve/main/config.json',
    'nghuyong/ernie-3.0-medium-zh': 'https://huggingface.co/nghuyong/ernie-3.0-medium-zh/resolve/main/config.json',
    'nghuyong/ernie-3.0-mini-zh': 'https://huggingface.co/nghuyong/ernie-3.0-mini-zh/resolve/main/config.json',
    'nghuyong/ernie-3.0-micro-zh': 'https://huggingface.co/nghuyong/ernie-3.0-micro-zh/resolve/main/config.json',
    'nghuyong/ernie-3.0-nano-zh': 'https://huggingface.co/nghuyong/ernie-3.0-nano-zh/resolve/main/config.json',
    'nghuyong/ernie-gram-zh': 'https://huggingface.co/nghuyong/ernie-gram-zh/resolve/main/config.json',
    'nghuyong/ernie-health-zh': 'https://huggingface.co/nghuyong/ernie-health-zh/resolve/main/config.json',
}







class 	lowerCamelCase__      (      __magic_name__       ):





			'''simple docstring'''
			lowerCamelCase						      =	'''ernie'''



			def __init__(							self  ,  __UpperCAmelCase=3_05_22  ,  __UpperCAmelCase=7_68  ,  __UpperCAmelCase=12  ,  __UpperCAmelCase=12  ,  __UpperCAmelCase=30_72  ,  __UpperCAmelCase="gelu"  ,  __UpperCAmelCase=0.1  ,  __UpperCAmelCase=0.1  ,  __UpperCAmelCase=5_12  ,  __UpperCAmelCase=2  ,  __UpperCAmelCase=3  ,  __UpperCAmelCase=False  ,  __UpperCAmelCase=0.0_2  ,  __UpperCAmelCase=1e-12  ,  __UpperCAmelCase=0  ,  __UpperCAmelCase="absolute"  ,  __UpperCAmelCase=True  ,  __UpperCAmelCase=None  ,  **__UpperCAmelCase  ,  )				->  str:
				super().__init__(pad_token_id=__UpperCAmelCase  ,  **__UpperCAmelCase						)

				_lowerCAmelCase			=vocab_size
				_lowerCAmelCase			=hidden_size
				_lowerCAmelCase			=num_hidden_layers
				_lowerCAmelCase			=num_attention_heads
				_lowerCAmelCase			=hidden_act
				_lowerCAmelCase			=intermediate_size
				_lowerCAmelCase			=hidden_dropout_prob
				_lowerCAmelCase			=attention_probs_dropout_prob
				_lowerCAmelCase			=max_position_embeddings
				_lowerCAmelCase			=type_vocab_size
				_lowerCAmelCase			=task_type_vocab_size
				_lowerCAmelCase			=use_task_id
				_lowerCAmelCase			=initializer_range
				_lowerCAmelCase			=layer_norm_eps
				_lowerCAmelCase			=position_embedding_type
				_lowerCAmelCase			=use_cache
				_lowerCAmelCase			=classifier_dropout







class 	lowerCamelCase__      (      __magic_name__       ):





			'''simple docstring'''



			@property
			def       _lowerCAmelCase	(							self						)				->  Mapping[str, Mapping[int, str]]:
				if self.task == "multiple-choice":
					_lowerCAmelCase			={0: """batch""", 1: """choice""", 2: """sequence"""}
				else:
					_lowerCAmelCase			={0: """batch""", 1: """sequence"""}
				return OrderedDict(
				    [
				        ("""input_ids""", dynamic_axis),
				        ("""attention_mask""", dynamic_axis),
				        ("""token_type_ids""", dynamic_axis),
				        ("""task_type_ids""", dynamic_axis),
				    ]						)



