





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


    """simple docstring"""


    UpperCamelCase   :  Optional[int]    =     None
    UpperCamelCase   :  Optional[Any]    =     ()
    UpperCamelCase   :  int    =     ()
    UpperCamelCase   :  List[Any]    =     True
    UpperCamelCase   :  str    =     True
    UpperCamelCase   :  int    =     True
    UpperCamelCase   :  Optional[int]    =     False



    def  UpperCAmelCase_		(      self					, _A					, _A					, _A=False					):
         __A     :						Dict          =    copy.deepcopy(_A					)

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



    @slow
    def  UpperCAmelCase_		(      self					):
         __A   ,    __A     :						Dict          =    self.model_tester.prepare_config_and_inputs_for_common()
         for model_class in self.all_model_classes[:2]:
              __A     :						List[str]          =    self._prepare_for_class(_A					, _A					)
              __A     :						str          =    model_class(_A					)

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

              __A     :						Optional[Any]          =    run_in_graph_mode()
              self.assertIsNotNone(_A					)



    @slow
    def  UpperCAmelCase_		(      self					):
         __A   ,    __A     :						List[str]          =    self.model_tester.prepare_config_and_inputs_for_common()
         for model_class in self.all_model_classes[:2]:
              __A     :						str          =    self._prepare_for_class(_A					, _A					)
              __A     :						int          =    model_class(_A					)

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

              __A     :						Any          =    run_in_graph_mode()
              self.assertIsNotNone(_A					)



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

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

                   __A     :						int          =    history.history['loss'][0]
                   self.assertTrue(not isnan(_A					)					)
                   __A     :						List[Any]          =    history.history['val_loss'][0]
                   self.assertTrue(not isnan(_A					)					)



    @slow
    def  UpperCAmelCase_		(      self					):
         __A   ,    __A     :						List[Any]          =    self.model_tester.prepare_config_and_inputs_for_common()
         __A     :						Dict          =    True
         __A     :						List[Any]          =    True

         if hasattr(_A					, 'use_cache'					):
              __A     :						List[Any]          =    True

         __A     :						int          =    getattr(self.model_tester					, 'encoder_seq_length'					, self.model_tester.seq_length					)
         __A     :						int          =    getattr(self.model_tester					, 'key_length'					, _A					)

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

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

                   self.assertEqual(len(_A					)					, _A					)

                   __A     :						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  UpperCAmelCase_		(      self					):
         tf.keras.mixed_precision.set_global_policy('mixed_float16'					)

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

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



    @slow
    def  UpperCAmelCase_		(      self					):
         __A   ,    __A     :						Optional[Any]          =    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"]
         __A     :						Union[str, Any]          =    {
             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
                   __A     :						Tuple          =    TFSharedEmbeddings(self.model_tester.vocab_size					, self.model_tester.hidden_size					, name='shared'					)
                   __A     :						List[Any]          =    False
                   __A     :						List[str]          =    main_layer_class(_A					, embed_tokens=_A					)
              else:
                   __A     :						List[Any]          =    main_layer_class(_A					)

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

              if hasattr(self.model_tester					, 'num_labels'					):
                   __A     :						Union[str, Any]          =    self.model_tester.num_labels
              else:
                   __A     :						Union[str, Any]          =    2

              __A     :						Union[str, Any]          =    tf.data.Dataset.from_tensor_slices(
                  (inputs_dict, np.ones((self.model_tester.batch_size, self.model_tester.seq_length, num_labels, 1)					))					).batch(1					)

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

         for model_class in self.all_model_classes[:2]:
              __A     :						Tuple          =    model_class(_A					)

              __A     :						Optional[Any]          =    copy.deepcopy(_A					)

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

              if not self.is_encoder_decoder:
                   __A     :						Dict          =    model.get_input_embeddings()(_A					)
              else:
                   __A     :						Dict          =    model.get_input_embeddings()(_A					)
                   __A     :						Dict          =    model.get_input_embeddings()(_A					)

              __A     :						int          =    self._prepare_for_class(_A					, _A					)

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

              __A     :						Optional[Any]          =    run_in_graph_mode()
              self.assertIsNotNone(_A					)



    def  UpperCAmelCase_		(      self					, _A					, _A					):
         __A     :						Optional[int]          =    []
         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
         __A     :						Tuple          =    []
         while len(_A					) < num_bad_tokens:
              __A     :						List[str]          =    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  UpperCAmelCase_		(      self					, _A					):







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





    def  UpperCAmelCase_		(      self					, _A					, _A					):







         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











from collections import OrderedDict
from typing import Mapping

from packaging import version

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


UpperCAmelCase						:       str        =					logging.get_logger(__name__)

UpperCAmelCase						:       Optional[Any]        =					{
    '''MBZUAI/swiftformer-xs''': '''https://huggingface.co/MBZUAI/swiftformer-xs/resolve/main/config.json''',
}





class 						_A( snake_case__       ):


				"""simple docstring"""


				UpperCamelCase   :  Union[str, Any]    =     '''swiftformer'''



				def __init__(      self					, _A=3					, _A=[3, 3, 6, 4]					, _A=[48, 56, 112, 220]					, _A=4					, _A=[True, True, True, True]					, _A="gelu"					, _A=3					, _A=2					, _A=1					, _A=0.0					, _A=True					, _A=1e-5					, _A=1e-5					, **_A					, ):
									super().__init__(**_A					)
									__A     :						List[Any]          =    num_channels
									__A     :						List[Any]          =    depths
									__A     :						int          =    embed_dims
									__A     :						Union[str, Any]          =    mlp_ratio
									__A     :						List[str]          =    downsamples
									__A     :						str          =    hidden_act
									__A     :						Any          =    down_patch_size
									__A     :						Any          =    down_stride
									__A     :						Any          =    down_pad
									__A     :						List[Any]          =    drop_path_rate
									__A     :						str          =    use_layer_scale
									__A     :						Dict          =    layer_scale_init_value
									__A     :						Union[str, Any]          =    batch_norm_eps





class 						_A( snake_case__       ):


				"""simple docstring"""


				UpperCamelCase   :  Optional[Any]    =     version.parse('''1.11'''       )



				@property
				def  UpperCAmelCase_		(      self					):
									return OrderedDict(
									    [
									        ('pixel_values', {0: 'batch', 1: 'num_channels', 2: 'height', 3: 'width'}),
									    ]					)





				@property
				def  UpperCAmelCase_		(      self					):
									return 1e-4




