




from __future__ import annotations

import copy
import inspect
import json
import math
import os
import tempfile
import unittest
from importlib import import_module

import numpy as np

from transformers import ViTMAEConfig
from transformers.file_utils import cached_property, is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, require_vision, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import TFViTMAEForPreTraining, TFViTMAEModel


if is_vision_available():
  from PIL import Image

  from transformers import ViTImageProcessor


class        UpperCamelCase_					:

    '''simple docstring'''






    def __init__(    self   :			Optional[int]      ,				UpperCAmelCase__   :			Any      ,				UpperCAmelCase__   :			List[Any]=13      ,				UpperCAmelCase__   :			Optional[Any]=30      ,				UpperCAmelCase__   :			int=2      ,				UpperCAmelCase__   :			int=3      ,				UpperCAmelCase__   :			Any=True      ,				UpperCAmelCase__   :			Union[str, Any]=True      ,				UpperCAmelCase__   :			Union[str, Any]=32      ,				UpperCAmelCase__   :			List[str]=2      ,				UpperCAmelCase__   :			List[Any]=4      ,				UpperCAmelCase__   :			int=37      ,				UpperCAmelCase__   :			int="gelu"      ,				UpperCAmelCase__   :			Any=0.1      ,				UpperCAmelCase__   :			List[Any]=0.1      ,				UpperCAmelCase__   :			List[str]=10      ,				UpperCAmelCase__   :			Any=0.02      ,				UpperCAmelCase__   :			Optional[int]=3      ,				UpperCAmelCase__   :			int=0.6      ,				UpperCAmelCase__   :			Union[str, Any]=None      ,				)		->Tuple:





           '''simple docstring'''

           A__   				=					parent
           A__   				=					batch_size
           A__   				=					image_size
           A__   				=					patch_size
           A__   				=					num_channels
           A__   				=					is_training
           A__   				=					use_labels
           A__   				=					hidden_size
           A__   				=					num_hidden_layers
           A__   				=					num_attention_heads
           A__   				=					intermediate_size
           A__   				=					hidden_act
           A__   				=					hidden_dropout_prob
           A__   				=					attention_probs_dropout_prob
           A__   				=					type_sequence_label_size
           A__   				=					initializer_range
           A__   				=					mask_ratio
           A__   				=					scope

           # in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
           # (we add 1 for the [CLS] token)
           A__   				=					(image_size // patch_size) ** 2
           A__   				=					int(math.ceil((1 - mask_ratio) * (num_patches + 1)))


    def  SCREAMING_SNAKE_CASE							(    self   :			Optional[Any])		->Any:





           '''simple docstring'''

           A__   				=					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

           A__   				=					None
           if self.use_labels:
                  A__   				=					ids_tensor([self.batch_size]      ,				self.type_sequence_label_size)

           A__   				=					self.get_config()

           return config, pixel_values, labels


    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->Dict:





           '''simple docstring'''

           return ViTMAEConfig(
               image_size=self.image_size      ,				patch_size=self.patch_size      ,				num_channels=self.num_channels      ,				hidden_size=self.hidden_size      ,				num_hidden_layers=self.num_hidden_layers      ,				num_attention_heads=self.num_attention_heads      ,				intermediate_size=self.intermediate_size      ,				decoder_hidden_size=self.hidden_size      ,				decoder_num_hidden_layers=self.num_hidden_layers      ,				decoder_num_attention_heads=self.num_attention_heads      ,				decoder_intermediate_size=self.intermediate_size      ,				hidden_act=self.hidden_act      ,				hidden_dropout_prob=self.hidden_dropout_prob      ,				attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,				is_decoder=UpperCAmelCase__      ,				initializer_range=self.initializer_range      ,				mask_ratio=self.mask_ratio      ,				)


    def  SCREAMING_SNAKE_CASE							(    self   :			int      ,				UpperCAmelCase__   :			Dict      ,				UpperCAmelCase__   :			Union[str, Any]      ,				UpperCAmelCase__   :			Any)		->List[str]:





           '''simple docstring'''

           A__   				=					TFViTMAEModel(config=UpperCAmelCase__)
           A__   				=					model(UpperCAmelCase__      ,				training=UpperCAmelCase__)
           self.parent.assertEqual(result.last_hidden_state.shape      ,				(self.batch_size, self.seq_length, self.hidden_size))


    def  SCREAMING_SNAKE_CASE							(    self   :			Any      ,				UpperCAmelCase__   :			List[Any]      ,				UpperCAmelCase__   :			Union[str, Any]      ,				UpperCAmelCase__   :			Union[str, Any])		->Union[str, Any]:





           '''simple docstring'''

           A__   				=					TFViTMAEForPreTraining(UpperCAmelCase__)
           A__   				=					model(UpperCAmelCase__      ,				training=UpperCAmelCase__)
           # expected sequence length = num_patches
           A__   				=					(self.image_size // self.patch_size) ** 2
           A__   				=					self.patch_size**2 * self.num_channels
           self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, num_patches, expected_num_channels))

           # test greyscale images
           A__   				=					1
           A__   				=					TFViTMAEForPreTraining(UpperCAmelCase__)

           A__   				=					floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
           A__   				=					model(UpperCAmelCase__      ,				training=UpperCAmelCase__)
           A__   				=					self.patch_size**2
           self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, num_patches, expected_num_channels))





    def  SCREAMING_SNAKE_CASE							(    self   :			Dict)		->Union[str, Any]:





           '''simple docstring'''

           A__   				=					self.prepare_config_and_inputs()
           ((A__)				,							(A__)				,							(A__))   				=					config_and_inputs
           A__   				=					{'''pixel_values''': pixel_values}
           return config, inputs_dict


@require_tf
class        UpperCamelCase_					(				UpperCAmelCase__	,      UpperCAmelCase__	,      unittest.TestCase  ):

    '''simple docstring'''






    UpperCAmelCase__													=   (TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
    UpperCAmelCase__													=   {'''feature-extraction''': TFViTMAEModel} if is_tf_available() else {}

    UpperCAmelCase__													=   False
    UpperCAmelCase__													=   False
    UpperCAmelCase__													=   False
    UpperCAmelCase__													=   False


    def  SCREAMING_SNAKE_CASE							(    self   :			Any)		->Optional[int]:





           '''simple docstring'''

           A__   				=					TFViTMAEModelTester(self)
           A__   				=					ConfigTester(self      ,				config_class=UpperCAmelCase__      ,				has_text_modality=UpperCAmelCase__      ,				hidden_size=37)


    def  SCREAMING_SNAKE_CASE							(    self   :			Dict)		->int:





           '''simple docstring'''

           self.config_tester.run_common_tests()


    @unittest.skip(reason='''ViTMAE does not use inputs_embeds''')
    def  SCREAMING_SNAKE_CASE							(    self   :			Tuple)		->Tuple:





           '''simple docstring'''

           pass


    def  SCREAMING_SNAKE_CASE							(    self   :			Optional[int])		->Dict:





           '''simple docstring'''

           A__				,							A__   				=					self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  A__   				=					model_class(UpperCAmelCase__)
                  self.assertIsInstance(model.get_input_embeddings()      ,				(tf.keras.layers.Layer))
                  A__   				=					model.get_output_embeddings()
                  self.assertTrue(x is None or isinstance(UpperCAmelCase__      ,				tf.keras.layers.Layer))


    def  SCREAMING_SNAKE_CASE							(    self   :			Tuple)		->List[Any]:





           '''simple docstring'''

           A__				,							A__   				=					self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  A__   				=					model_class(UpperCAmelCase__)
                  A__   				=					inspect.signature(model.call)
                  # signature.parameters is an OrderedDict => so arg_names order is deterministic
                  A__   				=					[*signature.parameters.keys()]

                  A__   				=					['''pixel_values''']
                  self.assertListEqual(arg_names[:1]      ,				UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			List[str])		->Any:





           '''simple docstring'''

           A__   				=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			Tuple)		->List[str]:





           '''simple docstring'''

           A__   				=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_pretraining(*UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			List[Any])		->int:





           '''simple docstring'''

           np.random.seed(2)

           A__				,							A__   				=					self.model_tester.prepare_config_and_inputs_for_common()

           A__   				=					int((config.image_size // config.patch_size) ** 2)
           A__   				=					np.random.uniform(size=(self.model_tester.batch_size, num_patches))

           for model_class in self.all_model_classes:
                  A__   				=					model_class(UpperCAmelCase__)
                  A__   				=					self._prepare_for_class(UpperCAmelCase__      ,				UpperCAmelCase__)

                  A__   				=					model(UpperCAmelCase__      ,				noise=UpperCAmelCase__)

                  A__   				=					copy.deepcopy(self._prepare_for_class(UpperCAmelCase__      ,				UpperCAmelCase__))
                  A__   				=					model(**UpperCAmelCase__      ,				noise=UpperCAmelCase__)
                  A__   				=					outputs_dict[0].numpy()
                  A__   				=					outputs_keywords[0].numpy()

                  self.assertLess(np.sum(np.abs(output_dict - output_keywords))      ,				1e-6)


    def  SCREAMING_SNAKE_CASE							(    self   :			Dict)		->Dict:





           '''simple docstring'''

           np.random.seed(2)

           A__				,							A__   				=					self.model_tester.prepare_config_and_inputs_for_common()

           A__   				=					int((config.image_size // config.patch_size) ** 2)
           A__   				=					np.random.uniform(size=(self.model_tester.batch_size, num_patches))

           def prepare_numpy_arrays(UpperCAmelCase__   :			Dict):
                  A__   				=					{}
                  for k, v in inputs_dict.items():
                         if tf.is_tensor(UpperCAmelCase__):
                                A__   				=					v.numpy()
                         else:
                                A__   				=					np.array(UpperCAmelCase__)

                  return inputs_np_dict

           for model_class in self.all_model_classes:
                  A__   				=					model_class(UpperCAmelCase__)

                  A__   				=					self._prepare_for_class(UpperCAmelCase__      ,				UpperCAmelCase__)
                  A__   				=					prepare_numpy_arrays(UpperCAmelCase__)

                  A__   				=					model(UpperCAmelCase__      ,				noise=UpperCAmelCase__)
                  A__   				=					model(**UpperCAmelCase__      ,				noise=UpperCAmelCase__)
                  self.assert_outputs_same(UpperCAmelCase__      ,				UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			str      ,				UpperCAmelCase__   :			List[Any]      ,				UpperCAmelCase__   :			Tuple      ,				UpperCAmelCase__   :			str)		->Any:





           '''simple docstring'''

           np.random.seed(2)

           A__   				=					int((tf_model.config.image_size // tf_model.config.patch_size) ** 2)
           A__   				=					np.random.uniform(size=(self.model_tester.batch_size, num_patches))
           A__   				=					tf.constant(UpperCAmelCase__)

           # Add `noise` argument.
           # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
           A__   				=					tf_noise

           super().check_pt_tf_models(UpperCAmelCase__      ,				UpperCAmelCase__      ,				UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			List[Any])		->Optional[Any]:





           '''simple docstring'''

           np.random.seed(2)

           A__				,							A__   				=					self.model_tester.prepare_config_and_inputs_for_common()

           A__   				=					{
               module_member
               for model_class in self.all_model_classes
               for module in (import_module(model_class.__module__),)
               for module_member_name in dir(UpperCAmelCase__)
               if module_member_name.endswith('''MainLayer''')
               # This condition is required, since `modeling_tf_clip.py` has 3 classes whose names end with `MainLayer`.
               and module_member_name[: -len('''MainLayer''')] == model_class.__name__[: -len('''Model''')]
               for module_member in (getattr(UpperCAmelCase__      ,				UpperCAmelCase__),)
               if isinstance(UpperCAmelCase__      ,				UpperCAmelCase__)
               and tf.keras.layers.Layer in module_member.__bases__
               and getattr(UpperCAmelCase__      ,				'''_keras_serializable'''      ,				UpperCAmelCase__)
           }

           A__   				=					int((config.image_size // config.patch_size) ** 2)
           A__   				=					np.random.uniform(size=(self.model_tester.batch_size, num_patches))
           A__   				=					tf.convert_to_tensor(UpperCAmelCase__)
           inputs_dict.update({'''noise''': noise})

           for main_layer_class in tf_main_layer_classes:
                  A__   				=					main_layer_class(UpperCAmelCase__)

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

                  A__   				=					tf.keras.Model(UpperCAmelCase__      ,				outputs=main_layer(UpperCAmelCase__))
                  A__   				=					model(UpperCAmelCase__)

                  with tempfile.TemporaryDirectory() as tmpdirname:
                         A__   				=					os.path.join(UpperCAmelCase__      ,				'''keras_model.h5''')
                         model.save(UpperCAmelCase__)
                         A__   				=					tf.keras.models.load_model(
                             UpperCAmelCase__      ,				custom_objects={main_layer_class.__name__: main_layer_class})
                         assert isinstance(UpperCAmelCase__      ,				tf.keras.Model)
                         A__   				=					model(UpperCAmelCase__)
                         self.assert_outputs_same(UpperCAmelCase__      ,				UpperCAmelCase__)


    @slow
    def  SCREAMING_SNAKE_CASE							(    self   :			Optional[int])		->Dict:





           '''simple docstring'''

           np.random.seed(2)

           A__				,							A__   				=					self.model_tester.prepare_config_and_inputs_for_common()

           A__   				=					int((config.image_size // config.patch_size) ** 2)
           A__   				=					np.random.uniform(size=(self.model_tester.batch_size, num_patches))

           for model_class in self.all_model_classes:
                  A__   				=					model_class(UpperCAmelCase__)
                  A__   				=					self._prepare_for_class(UpperCAmelCase__      ,				UpperCAmelCase__)
                  A__   				=					model(UpperCAmelCase__      ,				noise=UpperCAmelCase__)

                  if model_class.__name__ == "TFViTMAEModel":
                         A__   				=					outputs.last_hidden_state.numpy()
                         A__   				=					0
                  else:
                         A__   				=					outputs.logits.numpy()
                         A__   				=					0

                  with tempfile.TemporaryDirectory() as tmpdirname:
                         model.save_pretrained(UpperCAmelCase__      ,				saved_model=UpperCAmelCase__)
                         A__   				=					model_class.from_pretrained(UpperCAmelCase__)
                         A__   				=					model(UpperCAmelCase__      ,				noise=UpperCAmelCase__)

                         if model_class.__name__ == "TFViTMAEModel":
                                A__   				=					after_outputs['''last_hidden_state'''].numpy()
                                A__   				=					0
                         else:
                                A__   				=					after_outputs['''logits'''].numpy()
                                A__   				=					0

                         A__   				=					np.amax(np.abs(out_a - out_a))
                         self.assertLessEqual(UpperCAmelCase__      ,				1e-5)


    def  SCREAMING_SNAKE_CASE							(    self   :			Tuple)		->str:





           '''simple docstring'''

           np.random.seed(2)

           A__				,							A__   				=					self.model_tester.prepare_config_and_inputs_for_common()

           A__   				=					int((config.image_size // config.patch_size) ** 2)
           A__   				=					np.random.uniform(size=(self.model_tester.batch_size, num_patches))

           for model_class in self.all_model_classes:
                  A__   				=					model_class(UpperCAmelCase__)
                  A__   				=					self._prepare_for_class(UpperCAmelCase__      ,				UpperCAmelCase__)

                  A__   				=					model(UpperCAmelCase__      ,				noise=UpperCAmelCase__)
                  A__   				=					model.get_config()
                  # make sure that returned config is jsonifiable, which is required by keras
                  json.dumps(UpperCAmelCase__)
                  A__   				=					model_class.from_config(model.get_config())
                  # make sure it also accepts a normal config
                  A__   				=					model_class.from_config(model.config)
                  A__   				=					new_model(UpperCAmelCase__)  # Build model
                  new_model.set_weights(model.get_weights())
                  A__   				=					new_model(UpperCAmelCase__      ,				noise=UpperCAmelCase__)

                  self.assert_outputs_same(UpperCAmelCase__      ,				UpperCAmelCase__)


    @unittest.skip(
        reason='''ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results.''')
    def  SCREAMING_SNAKE_CASE							(    self   :			List[Any])		->Optional[Any]:





           '''simple docstring'''

           pass


    @unittest.skip(reason='''ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load''')
    def  SCREAMING_SNAKE_CASE							(    self   :			Any)		->str:





           '''simple docstring'''

           pass





    @slow
    def  SCREAMING_SNAKE_CASE							(    self   :			Optional[int])		->List[str]:





           '''simple docstring'''

           A__   				=					TFViTMAEModel.from_pretrained('''google/vit-base-patch16-224''')
           self.assertIsNotNone(UpperCAmelCase__)







def SCREAMING_SNAKE_CASE       ( )     -> int:



       """simple docstring"""

       A__   				=					Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''		)
       return image


@require_tf
@require_vision
class        UpperCamelCase_					(				unittest.TestCase  ):

    '''simple docstring'''






    @cached_property
    def  SCREAMING_SNAKE_CASE							(    self   :			List[str])		->Dict:





           '''simple docstring'''

           return ViTImageProcessor.from_pretrained('''facebook/vit-mae-base''') if is_vision_available() else None





    @slow
    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->List[str]:





           '''simple docstring'''

           np.random.seed(2)

           A__   				=					TFViTMAEForPreTraining.from_pretrained('''facebook/vit-mae-base''')

           A__   				=					self.default_image_processor
           A__   				=					prepare_img()
           A__   				=					image_processor(images=UpperCAmelCase__      ,				return_tensors='''tf''')

           # prepare a noise vector that will be also used for testing the TF model
           # (this way we can ensure that the PT and TF models operate on the same inputs)
           A__   				=					ViTMAEConfig()
           A__   				=					int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2)
           A__   				=					np.random.uniform(size=(1, num_patches))

           # forward pass
           A__   				=					model(**UpperCAmelCase__      ,				noise=UpperCAmelCase__)

           # verify the logits
           A__   				=					tf.convert_to_tensor([1, 196, 768])
           self.assertEqual(outputs.logits.shape      ,				UpperCAmelCase__)

           A__   				=					tf.convert_to_tensor(
               [[-0.0548, -1.7023, -0.9325], [0.3721, -0.5670, -0.2233], [0.8235, -1.3878, -0.3524]])

           tf.debugging.assert_near(outputs.logits[0, :3, :3]      ,				UpperCAmelCase__      ,				atol=1e-4)








from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowerCamelCase			:	Tuple							 =     logging.get_logger(__name__)

_lowerCamelCase			:	Tuple							 =     {
    """tiiuae/falcon-40b""": """https://huggingface.co/tiiuae/falcon-40b/resolve/main/config.json""",
    """tiiuae/falcon-7b""": """https://huggingface.co/tiiuae/falcon-7b/resolve/main/config.json""",
}


class        UpperCamelCase_					(				UpperCAmelCase__  ):

    '''simple docstring'''






    UpperCAmelCase__													=   '''falcon'''
    UpperCAmelCase__													=   ['''past_key_values''']


    def __init__(    self   :			Optional[Any]      ,				UpperCAmelCase__   :			str=65_024      ,				UpperCAmelCase__   :			List[Any]=4_544      ,				UpperCAmelCase__   :			Optional[Any]=32      ,				UpperCAmelCase__   :			Dict=71      ,				UpperCAmelCase__   :			Any=1e-5      ,				UpperCAmelCase__   :			Optional[int]=0.02      ,				UpperCAmelCase__   :			Optional[int]=True      ,				UpperCAmelCase__   :			Union[str, Any]=0.0      ,				UpperCAmelCase__   :			List[str]=0.0      ,				UpperCAmelCase__   :			str=None      ,				UpperCAmelCase__   :			Dict=False      ,				UpperCAmelCase__   :			List[str]=False      ,				UpperCAmelCase__   :			Optional[Any]=True      ,				UpperCAmelCase__   :			Optional[Any]=True      ,				UpperCAmelCase__   :			Optional[Any]=False      ,				UpperCAmelCase__   :			Dict=11      ,				UpperCAmelCase__   :			str=11      ,				**UpperCAmelCase__   :			Dict      ,				)		->Any:





           '''simple docstring'''

           A__   				=					vocab_size
           # Backward compatibility with n_embed kwarg
           A__   				=					kwargs.pop('''n_embed'''      ,				UpperCAmelCase__)
           A__   				=					hidden_size if n_embed is None else n_embed
           A__   				=					num_hidden_layers
           A__   				=					num_attention_heads
           A__   				=					layer_norm_epsilon
           A__   				=					initializer_range
           A__   				=					use_cache
           A__   				=					hidden_dropout
           A__   				=					attention_dropout

           A__   				=					bos_token_id
           A__   				=					eos_token_id
           A__   				=					num_attention_heads if num_kv_heads is None else num_kv_heads
           A__   				=					alibi
           A__   				=					new_decoder_architecture
           A__   				=					multi_query  # Ignored when new_decoder_architecture is True
           A__   				=					parallel_attn
           A__   				=					bias

           super().__init__(bos_token_id=UpperCAmelCase__      ,				eos_token_id=UpperCAmelCase__      ,				**UpperCAmelCase__)


    @property
    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->int:





           '''simple docstring'''

           return self.hidden_size // self.num_attention_heads





    @property
    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->Dict:





           '''simple docstring'''

           return not self.alibi


