


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:
    __UpperCAmelCase			:			Tuple			      =					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)]
        )
        __UpperCAmelCase			:			List[Any]			      =					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  UpperCAmelCase_						:
       '''simple docstring'''
       __UpperCamelCase   :  Dict				=       None
       __UpperCamelCase   :  List[Any]				=       ()
       __UpperCamelCase   :  Union[str, Any]				=       ()
       __UpperCamelCase   :  int				=       True
       __UpperCamelCase   :  int				=       True
       __UpperCamelCase   :  int				=       True
       __UpperCamelCase   :  str				=       False






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE=False    ):




              """simple docstring"""




              UpperCamelCase     :						List[str] 						=		copy.deepcopy(__A    )

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

              if return_labels:
                     if model_class in get_values(__A    ):
                            UpperCamelCase     :						int 						=		tf.ones(self.model_tester.batch_size			,       dtype=tf.intaa    )
                     elif model_class in get_values(__A    ):
                            UpperCamelCase     :						Optional[Any] 						=		tf.zeros(self.model_tester.batch_size			,       dtype=tf.intaa    )
                            UpperCamelCase     :						Optional[int] 						=		tf.zeros(self.model_tester.batch_size			,       dtype=tf.intaa    )
                     elif model_class in [
                         *get_values(__A    ),
                         *get_values(__A    ),
                     ]:
                            UpperCamelCase     :						List[str] 						=		tf.zeros(self.model_tester.batch_size			,       dtype=tf.intaa    )
                     elif model_class in get_values(__A    ):
                            UpperCamelCase     :						Union[str, Any] 						=		tf.zeros(self.model_tester.batch_size			,       dtype=tf.intaa    )
                     elif model_class in [
                         *get_values(__A    ),
                         *get_values(__A    ),
                         *get_values(__A    ),
                         *get_values(__A    ),
                         *get_values(__A    ),
                     ]:
                            UpperCamelCase     :						str 						=		tf.zeros(
                                (self.model_tester.batch_size, self.model_tester.seq_length)			,       dtype=tf.intaa    )
              return inputs_dict






       @slow
       def 			_lowercase (  self    ):




              """simple docstring"""




              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(__A			,       __A    )
                     UpperCamelCase     :						int 						=		model_class(__A    )

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

                     UpperCamelCase     :						Union[str, Any] 						=		run_in_graph_mode()
                     self.assertIsNotNone(__A    )






       @slow
       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						Tuple 						=		self.model_tester.prepare_config_and_inputs_for_common()
              for model_class in self.all_model_classes[:2]:
                     UpperCamelCase     :						int 						=		self._prepare_for_class(__A			,       __A    )
                     UpperCamelCase     :						Tuple 						=		model_class(__A    )

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

                     UpperCamelCase     :						int 						=		run_in_graph_mode()
                     self.assertIsNotNone(__A    )






       @slow
       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						Union[str, Any] 						=		self.model_tester.prepare_config_and_inputs_for_common()
              for model_class in self.all_model_classes[:2]:
                     UpperCamelCase     :						int 						=		model_class(__A    )
                     if getattr(__A			,       '''hf_compute_loss'''			,       __A    ):
                            # Test that model correctly compute the loss with kwargs
                            UpperCamelCase     :						Optional[Any] 						=		self._prepare_for_class(inputs_dict.copy()			,       __A			,       return_labels=__A    )
                            # Is there a better way to remove these decoder inputs?
                            UpperCamelCase     :						Union[str, Any] 						=		{
                                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     :						Union[str, Any] 						=		{
                                """labels""",
                                """label""",
                                """label_ids""",
                                """start_positions""",
                                """start_position""",
                                """end_positions""",
                                """end_position""",
                                """next_sentence_label""",
                            }
                            UpperCamelCase     :						str 						=		possible_label_cols.intersection(set(__A    )    )
                            self.assertGreater(len(__A    )			,       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     :						int 						=		{key: val for key, val in prepared_for_class.items() if key not in label_names}
                            self.assertGreater(len(__A    )			,       0    )

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

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

                            UpperCamelCase     :						str 						=		history.history["""loss"""][0]
                            self.assertTrue(not isnan(__A    )    )
                            UpperCamelCase     :						Union[str, Any] 						=		history.history["""val_loss"""][0]
                            self.assertTrue(not isnan(__A    )    )






       @slow
       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						Tuple 						=		self.model_tester.prepare_config_and_inputs_for_common()
              UpperCamelCase     :						str 						=		True
              UpperCamelCase     :						Optional[int] 						=		True

              if hasattr(__A			,       '''use_cache'''    ):
                     UpperCamelCase     :						int 						=		True

              UpperCamelCase     :						Dict 						=		getattr(self.model_tester			,       '''encoder_seq_length'''			,       self.model_tester.seq_length    )
              UpperCamelCase     :						Optional[int] 						=		getattr(self.model_tester			,       '''key_length'''			,       __A    )

              for model_class in self.all_model_classes[:2]:
                     UpperCamelCase     :						Dict 						=		self._prepare_for_class(__A			,       __A    )
                     UpperCamelCase     :						int 						=		model_class(__A    )
                     model.build()
                     UpperCamelCase     :						Union[str, Any] 						=		len(model(__A    )    )

                     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     :						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(__A			,       saved_model=__A    )
                            UpperCamelCase     :						Any 						=		os.path.join(__A			,       '''saved_model'''			,       '''1'''    )
                            UpperCamelCase     :						Dict 						=		tf.keras.models.load_model(__A    )
                            UpperCamelCase     :						Tuple 						=		model(__A    )

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

                            self.assertEqual(len(__A    )			,       __A    )

                            UpperCamelCase     :						Union[str, Any] 						=		getattr(
                                self.model_tester			,       '''expected_num_hidden_layers'''			,       self.model_tester.num_hidden_layers + 1    )

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

                            self.assertEqual(len(__A    )			,       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    ):




              """simple docstring"""




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

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

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






       @slow
       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						Tuple 						=		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(__A    )
                  if module_member_name.endswith('''MainLayer'''    )
                  for module_member in (getattr(__A			,       __A    ),)
                  if isinstance(__A			,       __A    )
                  and tf.keras.layers.Layer in module_member.__bases__
                  and getattr(__A			,       '''_keras_serializable'''			,       __A    )
              }

              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     :						Any 						=		TFSharedEmbeddings(self.model_tester.vocab_size			,       self.model_tester.hidden_size			,       name='''shared'''    )
                            UpperCamelCase     :						Tuple 						=		False
                            UpperCamelCase     :						Dict 						=		main_layer_class(__A			,       embed_tokens=__A    )
                     else:
                            UpperCamelCase     :						Optional[int] 						=		main_layer_class(__A    )

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

                     UpperCamelCase     :						Optional[int] 						=		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     :						Any 						=		main_layer(__A    )[0]
                     UpperCamelCase     :						Optional[Any] 						=		tf.keras.layers.Dense(__A			,       activation='''softmax'''			,       name='''outputs'''    )(__A    )
                     UpperCamelCase     :						Any 						=		tf.keras.models.Model(inputs=__A			,       outputs=[outputs]    )

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

                     with tempfile.TemporaryDirectory() as tmpdirname:
                            UpperCamelCase     :						Union[str, Any] 						=		os.path.join(__A			,       '''keras_model.h5'''    )
                            model.save(__A    )
                            if "T5" in main_layer_class.__name__:
                                   UpperCamelCase     :						Union[str, Any] 						=		tf.keras.models.load_model(
                                       __A			,       custom_objects={
                                           main_layer_class.__name__: main_layer_class,
                                           '''TFSharedEmbeddings''': TFSharedEmbeddings,
                                       }			,       )
                            else:
                                   UpperCamelCase     :						Optional[Any] 						=		tf.keras.models.load_model(
                                       __A			,       custom_objects={main_layer_class.__name__: main_layer_class}    )
                            assert isinstance(__A			,       tf.keras.Model    )
                            model(__A    )






       @slow
       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						Any 						=		self.model_tester.prepare_config_and_inputs_for_common()

              for model_class in self.all_model_classes[:2]:
                     UpperCamelCase     :						Dict 						=		model_class(__A    )

                     UpperCamelCase     :						Union[str, Any] 						=		copy.deepcopy(__A    )

                     if not self.is_encoder_decoder:
                            UpperCamelCase     :						Optional[int] 						=		inputs["""input_ids"""]
                            del inputs["input_ids"]
                     else:
                            UpperCamelCase     :						Union[str, Any] 						=		inputs["""input_ids"""]
                            UpperCamelCase     :						List[Any] 						=		inputs.get('''decoder_input_ids'''			,       __A    )
                            del inputs["input_ids"]
                            inputs.pop('''decoder_input_ids'''			,       __A    )

                     if not self.is_encoder_decoder:
                            UpperCamelCase     :						str 						=		model.get_input_embeddings()(__A    )
                     else:
                            UpperCamelCase     :						List[str] 						=		model.get_input_embeddings()(__A    )
                            UpperCamelCase     :						Dict 						=		model.get_input_embeddings()(__A    )

                     UpperCamelCase     :						Optional[Any] 						=		self._prepare_for_class(__A			,       __A    )

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

                     UpperCamelCase     :						Union[str, Any] 						=		run_in_graph_mode()
                     self.assertIsNotNone(__A    )






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




              UpperCamelCase     :						Dict 						=		[]
              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     :						Any 						=		[]
              while len(__A    ) < num_bad_tokens:
                     UpperCamelCase     :						Any 						=		tf.squeeze(ids_tensor((1, 1)			,       self.model_tester.vocab_size    )			,       0    ).numpy()[0]
                     if token not in special_tokens:
                            bad_tokens.append(__A    )
              return bad_tokens






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE    ):




              """simple docstring"""




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






       def 			_lowercase (  self			,       __SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ):




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



"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCAmelCase 					=    {
    'configuration_marian': ['MARIAN_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MarianConfig', 'MarianOnnxConfig'],
}

try:
				if not is_sentencepiece_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__UpperCAmelCase 					=    ['MarianTokenizer']

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__UpperCAmelCase 					=    [
				    'MARIAN_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'MarianForCausalLM',
				    'MarianModel',
				    'MarianMTModel',
				    'MarianPreTrainedModel',
				]

try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__UpperCAmelCase 					=    ['TFMarianModel', 'TFMarianMTModel', 'TFMarianPreTrainedModel']

try:
				if not is_flax_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__UpperCAmelCase 					=    ['FlaxMarianModel', 'FlaxMarianMTModel', 'FlaxMarianPreTrainedModel']

if TYPE_CHECKING:
				from .configuration_marian import MARIAN_PRETRAINED_CONFIG_ARCHIVE_MAP, MarianConfig, MarianOnnxConfig

				try:
								if not is_sentencepiece_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_marian import MarianTokenizer

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_marian import (
								    MARIAN_PRETRAINED_MODEL_ARCHIVE_LIST,
								    MarianForCausalLM,
								    MarianModel,
								    MarianMTModel,
								    MarianPreTrainedModel,
								)

				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_marian import TFMarianModel, TFMarianMTModel, TFMarianPreTrainedModel

				try:
								if not is_flax_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_flax_marian import FlaxMarianModel, FlaxMarianMTModel, FlaxMarianPreTrainedModel

else:
				import sys

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






