
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 			lowercase__:



	"""simple docstring"""




	a      :Dict   					=     None
	a      :Union[str, Any]   					=     ()
	a      :List[Any]   					=     ()
	a      :List[Any]   					=     True
	a      :Optional[int]   					=     True
	a      :Tuple   					=     True
	a      :Union[str, Any]   					=     False
	def 						_lowercase  (	self  :      str	,							SCREAMING_SNAKE_CASE_  :      Optional[int]	,							SCREAMING_SNAKE_CASE_  :      List[Any]	,							SCREAMING_SNAKE_CASE_  :      int=False	)	->					Tuple:
						lowercase_					  = copy.deepcopy(__lowerCamelCase	)

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

						if return_labels:
											if model_class in get_values(__lowerCamelCase	):
																lowercase_					  = tf.ones(self.model_tester.batch_size	,							dtype=tf.intaa	)
											elif model_class in get_values(__lowerCamelCase	):
																lowercase_					  = tf.zeros(self.model_tester.batch_size	,							dtype=tf.intaa	)
																lowercase_					  = tf.zeros(self.model_tester.batch_size	,							dtype=tf.intaa	)
											elif model_class in [
											    *get_values(__lowerCamelCase	),
											    *get_values(__lowerCamelCase	),
											]:
																lowercase_					  = tf.zeros(self.model_tester.batch_size	,							dtype=tf.intaa	)
											elif model_class in get_values(__lowerCamelCase	):
																lowercase_					  = tf.zeros(self.model_tester.batch_size	,							dtype=tf.intaa	)
											elif model_class in [
											    *get_values(__lowerCamelCase	),
											    *get_values(__lowerCamelCase	),
											    *get_values(__lowerCamelCase	),
											    *get_values(__lowerCamelCase	),
											    *get_values(__lowerCamelCase	),
											]:
																lowercase_					  = tf.zeros(
																    (self.model_tester.batch_size, self.model_tester.seq_length)	,							dtype=tf.intaa	)
						return inputs_dict
	@slow
	def 						_lowercase  (	self  :      Dict	)	->					Dict:
						lowercase_					  = self.model_tester.prepare_config_and_inputs_for_common()
						for model_class in self.all_model_classes[:2]:
											lowercase_					  = self._prepare_for_class(__lowerCamelCase	,							__lowerCamelCase	)
											lowercase_					  = model_class(__lowerCamelCase	)

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

											lowercase_					  = run_in_graph_mode()
											self.assertIsNotNone(__lowerCamelCase	)
	@slow
	def 						_lowercase  (	self  :      Optional[Any]	)	->					List[str]:
						lowercase_					  = self.model_tester.prepare_config_and_inputs_for_common()
						for model_class in self.all_model_classes[:2]:
											lowercase_					  = self._prepare_for_class(__lowerCamelCase	,							__lowerCamelCase	)
											lowercase_					  = model_class(__lowerCamelCase	)

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

											lowercase_					  = run_in_graph_mode()
											self.assertIsNotNone(__lowerCamelCase	)
	@slow
	def 						_lowercase  (	self  :      Optional[Any]	)	->					Optional[Any]:
						lowercase_					  = self.model_tester.prepare_config_and_inputs_for_common()
						for model_class in self.all_model_classes[:2]:
											lowercase_					  = model_class(__lowerCamelCase	)
											if getattr(__lowerCamelCase	,							'''hf_compute_loss'''	,							__lowerCamelCase	):
																# Test that model correctly compute the loss with kwargs
																lowercase_					  = self._prepare_for_class(inputs_dict.copy()	,							__lowerCamelCase	,							return_labels=__lowerCamelCase	)
																# Is there a better way to remove these decoder inputs?
																lowercase_					  = {
																    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")
																}

																lowercase_					  = {
																    "labels",
																    "label",
																    "label_ids",
																    "start_positions",
																    "start_position",
																    "end_positions",
																    "end_position",
																    "next_sentence_label",
																}
																lowercase_					  = possible_label_cols.intersection(set(__lowerCamelCase	)	)
																self.assertGreater(len(__lowerCamelCase	)	,							0	,							msg='''No matching label names found!'''	)
																lowercase_					  = {key: val for key, val in prepared_for_class.items() if key in label_names}
																lowercase_					  = {key: val for key, val in prepared_for_class.items() if key not in label_names}
																self.assertGreater(len(__lowerCamelCase	)	,							0	)

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

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

																lowercase_					  = history.history["loss"][0]
																self.assertTrue(not isnan(__lowerCamelCase	)	)
																lowercase_					  = history.history["val_loss"][0]
																self.assertTrue(not isnan(__lowerCamelCase	)	)
	@slow
	def 						_lowercase  (	self  :      Union[str, Any]	)	->					Optional[Any]:
						lowercase_					  = self.model_tester.prepare_config_and_inputs_for_common()
						lowercase_					  = True
						lowercase_					  = True

						if hasattr(__lowerCamelCase	,							'''use_cache'''	):
											lowercase_					  = True

						lowercase_					  = getattr(self.model_tester	,							'''encoder_seq_length'''	,							self.model_tester.seq_length	)
						lowercase_					  = getattr(self.model_tester	,							'''key_length'''	,							__lowerCamelCase	)

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

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

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

																self.assertEqual(len(__lowerCamelCase	)	,							__lowerCamelCase	)

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

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

																self.assertEqual(len(__lowerCamelCase	)	,							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  :      Tuple	)	->					int:
						tf.keras.mixed_precision.set_global_policy('''mixed_float16'''	)

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

																self.assertIsNotNone(__lowerCamelCase	)
						finally:
											tf.keras.mixed_precision.set_global_policy('''float32'''	)
	@slow
	def 						_lowercase  (	self  :      Optional[Any]	)	->					Optional[int]:
						lowercase_					  = 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"]
						lowercase_					  = {
						    module_member
						    for model_class in self.all_model_classes
						    for module in (import_module(model_class.__module__	),)
						    for module_member_name in dir(__lowerCamelCase	)
						    if module_member_name.endswith('''MainLayer'''	)
						    for module_member in (getattr(__lowerCamelCase	,							__lowerCamelCase	),)
						    if isinstance(__lowerCamelCase	,							__lowerCamelCase	)
						    and tf.keras.layers.Layer in module_member.__bases__
						    and getattr(__lowerCamelCase	,							'''_keras_serializable'''	,							__lowerCamelCase	)
						}

						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
																lowercase_					  = TFSharedEmbeddings(self.model_tester.vocab_size	,							self.model_tester.hidden_size	,							name='''shared'''	)
																lowercase_					  = False
																lowercase_					  = main_layer_class(__lowerCamelCase	,							embed_tokens=__lowerCamelCase	)
											else:
																lowercase_					  = main_layer_class(__lowerCamelCase	)

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

											if hasattr(self.model_tester	,							'''num_labels'''	):
																lowercase_					  = self.model_tester.num_labels
											else:
																lowercase_					  = 2

											lowercase_					  = 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	)

											lowercase_					  = main_layer(__lowerCamelCase	)[0]
											lowercase_					  = tf.keras.layers.Dense(__lowerCamelCase	,							activation='''softmax'''	,							name='''outputs'''	)(__lowerCamelCase	)
											lowercase_					  = tf.keras.models.Model(inputs=__lowerCamelCase	,							outputs=[outputs]	)

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

											with tempfile.TemporaryDirectory() as tmpdirname:
																lowercase_					  = os.path.join(__lowerCamelCase	,							'''keras_model.h5'''	)
																model.save(__lowerCamelCase	)
																if "T5" in main_layer_class.__name__:
																					lowercase_					  = tf.keras.models.load_model(
																					    __lowerCamelCase	,							custom_objects={
																					        main_layer_class.__name__: main_layer_class,
																					        '''TFSharedEmbeddings''': TFSharedEmbeddings,
																					    }	,							)
																else:
																					lowercase_					  = tf.keras.models.load_model(
																					    __lowerCamelCase	,							custom_objects={main_layer_class.__name__: main_layer_class}	)
																assert isinstance(__lowerCamelCase	,							tf.keras.Model	)
																model(__lowerCamelCase	)
	@slow
	def 						_lowercase  (	self  :      Any	)	->					Tuple:
						lowercase_					  = self.model_tester.prepare_config_and_inputs_for_common()

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

											lowercase_					  = copy.deepcopy(__lowerCamelCase	)

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

											if not self.is_encoder_decoder:
																lowercase_					  = model.get_input_embeddings()(__lowerCamelCase	)
											else:
																lowercase_					  = model.get_input_embeddings()(__lowerCamelCase	)
																lowercase_					  = model.get_input_embeddings()(__lowerCamelCase	)

											lowercase_					  = self._prepare_for_class(__lowerCamelCase	,							__lowerCamelCase	)

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

											lowercase_					  = run_in_graph_mode()
											self.assertIsNotNone(__lowerCamelCase	)
	def 						_lowercase  (	self  :      Union[str, Any]	,							SCREAMING_SNAKE_CASE_  :      List[str]	,							SCREAMING_SNAKE_CASE_  :      Union[str, Any]	)	->					List[str]:
						lowercase_					  = []
						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
						lowercase_					  = []
						while len(__lowerCamelCase	) < num_bad_tokens:
											lowercase_					  = tf.squeeze(ids_tensor((1, 1)	,							self.model_tester.vocab_size	)	,							0	).numpy()[0]
											if token not in special_tokens:
																bad_tokens.append(__lowerCamelCase	)
						return bad_tokens
	def 						_lowercase  (	self  :      str	,							SCREAMING_SNAKE_CASE_  :      Optional[Any]	)	->					Optional[int]:


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




	def 						_lowercase  (	self  :      str	,							SCREAMING_SNAKE_CASE_  :      Optional[int]	,							SCREAMING_SNAKE_CASE_  :      Any	)	->					List[Any]:


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






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


__a        = {
    'configuration_donut_swin': ['DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP', 'DonutSwinConfig'],
    'processing_donut': ['DonutProcessor'],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__a        = [
						    'DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST',
						    'DonutSwinModel',
						    'DonutSwinPreTrainedModel',
						]

try:
						if not is_vision_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__a        = ['DonutFeatureExtractor']
						__a        = ['DonutImageProcessor']


if TYPE_CHECKING:
						from .configuration_donut_swin import DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, DonutSwinConfig
						from .processing_donut import DonutProcessor

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_donut_swin import (
												    DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
												    DonutSwinModel,
												    DonutSwinPreTrainedModel,
												)

						try:
												if not is_vision_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .feature_extraction_donut import DonutFeatureExtractor
												from .image_processing_donut import DonutImageProcessor

else:
						import sys

						__a        = _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




