




import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import AutoConfig


__SCREAMING_SNAKE_CASE     : List[str]      =logging.get_logger(__name__)






class       A      (		_lowerCAmelCase		):
     """simple docstring"""

     _snake_case:						Union[str, Any]  ="""speech-encoder-decoder"""
     _snake_case:						Union[str, Any]  =True




     def __init__(     self	: Tuple						,	**A__	: Optional[int]      )					->							Optional[int]:
           super().__init__(**A__      )
           if "encoder" not in kwargs or "decoder" not in kwargs:
                 raise ValueError(
                     F"""A configuraton of type {self.model_type} cannot be instantiated because not both `encoder` and"""
                     F""" `decoder` sub-configurations are passed, but only {kwargs}"""      )

           lowerCamelCase__   	=					kwargs.pop('encoder'      )
           lowerCamelCase__   	=					encoder_config.pop('model_type'      )
           lowerCamelCase__   	=					kwargs.pop('decoder'      )
           lowerCamelCase__   	=					decoder_config.pop('model_type'      )

           lowerCamelCase__   	=					AutoConfig.for_model(A__						,	**A__      )
           lowerCamelCase__   	=					AutoConfig.for_model(A__						,	**A__      )
           lowerCamelCase__   	=					True




     @classmethod
     def 			__snake_case (     cls	: Any						,	A__	: Union[str, Any]						,	A__	: Union[str, Any]						,	**A__	: int      )					->							List[str]:
           logger.info('Setting `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config'      )
           lowerCamelCase__   	=					True
           lowerCamelCase__   	=					True

           return cls(encoder=encoder_config.to_dict()						,	decoder=decoder_config.to_dict()						,	**A__      )




     def 			__snake_case (     self	: Any      )					->							Any:
           lowerCamelCase__   	=					copy.deepcopy(self.__dict__      )
           lowerCamelCase__   	=					self.encoder.to_dict()
           lowerCamelCase__   	=					self.decoder.to_dict()
           lowerCamelCase__   	=					self.__class__.model_type
           return output
"""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__       =	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__       =	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  :     List[str]      	=    None
				_lowerCAmelCase  :     List[Any]      	=    ()
				_lowerCAmelCase  :     Union[str, Any]      	=    ()
				_lowerCAmelCase  :     List[Any]      	=    True
				_lowerCAmelCase  :     Dict      	=    True
				_lowerCAmelCase  :     Any      	=    True
				_lowerCAmelCase  :     Any      	=    False

				def 		snake_case			(					self		,     lowerCAmelCase		,     lowerCAmelCase		,     lowerCAmelCase=False     ):






										"""simple docstring"""

										snake_case             =      copy.deepcopy(lowerCAmelCase     )

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

				@slow
				def 		snake_case			(					self     ):






										"""simple docstring"""

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

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

																snake_case             =      run_in_graph_mode()
																self.assertIsNotNone(lowerCAmelCase     )

				@slow
				def 		snake_case			(					self     ):






										"""simple docstring"""

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

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

																snake_case             =      run_in_graph_mode()
																self.assertIsNotNone(lowerCAmelCase     )

				@slow
				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case       ,snake_case             =      self.model_tester.prepare_config_and_inputs_for_common()
										for model_class in self.all_model_classes[:2]:
																snake_case             =      model_class(lowerCAmelCase     )
																if getattr(lowerCAmelCase		,     'hf_compute_loss'		,     lowerCAmelCase     ):
																						# Test that model correctly compute the loss with kwargs
																						snake_case             =      self._prepare_for_class(inputs_dict.copy()		,     lowerCAmelCase		,     return_labels=lowerCAmelCase     )
																						# Is there a better way to remove these decoder inputs?
																						snake_case             =      {
																						    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             =      {
																						    'labels',
																						    'label',
																						    'label_ids',
																						    'start_positions',
																						    'start_position',
																						    'end_positions',
																						    'end_position',
																						    'next_sentence_label',
																						}
																						snake_case             =      possible_label_cols.intersection(set(lowerCAmelCase     )     )
																						self.assertGreater(len(lowerCAmelCase     )		,     0		,     msg='No matching label names found!'     )
																						snake_case             =      {key: val for key, val in prepared_for_class.items() if key in label_names}
																						snake_case             =      {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
																						snake_case             =      model.fit(
																						    lowerCAmelCase		,     validation_data=lowerCAmelCase		,     steps_per_epoch=1		,     validation_steps=1		,     shuffle=lowerCAmelCase		,     verbose=0		,     )
																						snake_case             =      history.history['loss'][0]
																						self.assertTrue(not isnan(lowerCAmelCase     )     )
																						snake_case             =      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.
																						snake_case             =      model_class(lowerCAmelCase     )
																						model.compile(optimizer=tf.keras.optimizers.SGD(0.0     )		,     jit_compile=lowerCAmelCase     )
																						snake_case             =      model.fit(
																						    lowerCAmelCase		,     lowerCAmelCase		,     validation_data=(inputs_minus_labels, labels)		,     steps_per_epoch=1		,     validation_steps=1		,     shuffle=lowerCAmelCase		,     verbose=0		,     )

																						snake_case             =      history.history['loss'][0]
																						self.assertTrue(not isnan(lowerCAmelCase     )     )
																						snake_case             =      history.history['val_loss'][0]
																						self.assertTrue(not isnan(lowerCAmelCase     )     )

				@slow
				def 		snake_case			(					self     ):






										"""simple docstring"""

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

										if hasattr(lowerCAmelCase		,     'use_cache'     ):
																snake_case             =      True

										snake_case             =      getattr(self.model_tester		,     'encoder_seq_length'		,     self.model_tester.seq_length     )
										snake_case             =      getattr(self.model_tester		,     'key_length'		,     lowerCAmelCase     )

										for model_class in self.all_model_classes[:2]:
																snake_case             =      self._prepare_for_class(lowerCAmelCase		,     lowerCAmelCase     )
																snake_case             =      model_class(lowerCAmelCase     )
																model.build()
																snake_case             =      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:
																												snake_case             =      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     )
																						snake_case             =      os.path.join(lowerCAmelCase		,     'saved_model'		,     '1'     )
																						snake_case             =      tf.keras.models.load_model(lowerCAmelCase     )
																						snake_case             =      model(lowerCAmelCase     )

																						if self.is_encoder_decoder:
																												snake_case             =      outputs['encoder_hidden_states']
																												snake_case             =      outputs['encoder_attentions']
																						else:
																												snake_case             =      outputs['hidden_states']
																												snake_case             =      outputs['attentions']

																						self.assertEqual(len(lowerCAmelCase     )		,     lowerCAmelCase     )

																						snake_case             =      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 		snake_case			(					self     ):






										"""simple docstring"""

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

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

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

				@slow
				def 		snake_case			(					self     ):






										"""simple docstring"""

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

																snake_case             =      {
																    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             =      self.model_tester.num_labels
																else:
																						snake_case             =      2

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






										"""simple docstring"""

										snake_case       ,snake_case             =      self.model_tester.prepare_config_and_inputs_for_common()

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

																snake_case             =      copy.deepcopy(lowerCAmelCase     )

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

																if not self.is_encoder_decoder:
																						snake_case             =      model.get_input_embeddings()(lowerCAmelCase     )
																else:
																						snake_case             =      model.get_input_embeddings()(lowerCAmelCase     )
																						snake_case             =      model.get_input_embeddings()(lowerCAmelCase     )

																snake_case             =      self._prepare_for_class(lowerCAmelCase		,     lowerCAmelCase     )

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

																snake_case             =      run_in_graph_mode()
																self.assertIsNotNone(lowerCAmelCase     )

				def 		snake_case			(					self		,     lowerCAmelCase		,     lowerCAmelCase     ):






										"""simple docstring"""

										snake_case             =      []
										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             =      []
										while len(lowerCAmelCase     ) < num_bad_tokens:
																snake_case             =      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 		snake_case			(					self		,     lowerCAmelCase     ):






										"""simple docstring"""

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

				def 		snake_case			(					self		,     lowerCAmelCase		,     lowerCAmelCase     ):






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


