






"""simple docstring"""




import inspect
import math
import tempfile
import unittest

import numpy as np

from transformers import ViTMAEConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch
     from torch import nn

     from transformers import ViTMAEForPreTraining, ViTMAEModel
     from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import ViTImageProcessor




class    lowercase_    :






      def __init__(    self     ,     a_     ,     a_=1_3     ,     a_=3_0     ,     a_=2     ,     a_=3     ,     a_=True     ,     a_=True     ,     a_=3_2     ,     a_=5     ,     a_=4     ,     a_=3_7     ,     a_="gelu"     ,     a_=0.1     ,     a_=0.1     ,     a_=1_0     ,     a_=0.02     ,     a_=3     ,     a_=0.6     ,     a_=None     ,     )	->List[str]:







        '''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 			lowerCamelCase__    (    self	)	->List[str]:







        '''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 			lowerCamelCase__    (    self	)	->str:







        '''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     ,     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 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_	)	->Union[str, Any]:







        '''simple docstring'''


        _a							=	ViTMAEModel(config=__UpperCamelCase	)
        model.to(__UpperCamelCase	)
        model.eval()
        _a							=	model(__UpperCamelCase	)
        self.parent.assertEqual(result.last_hidden_state.shape     ,     (self.batch_size, self.seq_length, self.hidden_size)	)






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_	)	->Optional[Any]:







        '''simple docstring'''


        _a							=	ViTMAEForPreTraining(__UpperCamelCase	)
        model.to(__UpperCamelCase	)
        model.eval()
        _a							=	model(__UpperCamelCase	)
        _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							=	ViTMAEForPreTraining(__UpperCamelCase	)
        model.to(__UpperCamelCase	)
        model.eval()
        _a							=	floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
        _a							=	model(__UpperCamelCase	)
        _a							=	self.patch_size**2
        self.parent.assertEqual(result.logits.shape     ,     (self.batch_size, num_patches, expected_num_channels)	)






      def 			lowerCamelCase__    (    self	)	->Optional[int]:







        '''simple docstring'''


        _a							=	self.prepare_config_and_inputs()
        _a			,	_a			,	_a							=	config_and_inputs
        _a							=	{"pixel_values": pixel_values}
        return config, inputs_dict







@require_torch
class    lowercase_    (_UpperCAmelCase,    _UpperCAmelCase,    unittest.TestCase     ):
      A__	:		Any       			=				(ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
      A__	:		str       			=				{'''feature-extraction''': ViTMAEModel} if is_torch_available() else {}

      A__	:		List[str]       			=				False
      A__	:		Any       			=				False
      A__	:		List[str]       			=				False
      A__	:		int       			=				False






      def 			lowerCamelCase__    (    self	)	->List[str]:







        '''simple docstring'''


        _a							=	ViTMAEModelTester(self	)
        _a							=	ConfigTester(self     ,     config_class=__UpperCamelCase     ,     has_text_modality=__UpperCamelCase     ,     hidden_size=3_7	)






      def 			lowerCamelCase__    (    self	)	->str:







        '''simple docstring'''


        self.config_tester.run_common_tests()






      @unittest.skip(reason="ViTMAE does not use inputs_embeds"	)
      def 			lowerCamelCase__    (    self	)	->Union[str, Any]:







        '''simple docstring'''


        pass






      def 			lowerCamelCase__    (    self	)	->str:







        '''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()     ,     (nn.Module)	)
          _a							=	model.get_output_embeddings()
          self.assertTrue(x is None or isinstance(__UpperCamelCase     ,     nn.Linear	)	)






      def 			lowerCamelCase__    (    self	)	->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.forward	)
          # 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 			lowerCamelCase__    (    self	)	->List[Any]:







        '''simple docstring'''


        _a							=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__UpperCamelCase	)






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







        '''simple docstring'''


        _a							=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*__UpperCamelCase	)






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_	)	->Union[str, Any]:







        '''simple docstring'''


        np.random.seed(2	)

        _a							=	int((pt_model.config.image_size // pt_model.config.patch_size) ** 2	)
        _a							=	np.random.uniform(size=(self.model_tester.batch_size, num_patches)	)
        _a							=	torch.from_numpy(__UpperCamelCase	)

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

        super().check_pt_tf_models(__UpperCamelCase     ,     __UpperCamelCase     ,     __UpperCamelCase	)






      def 			lowerCamelCase__    (    self	)	->Tuple:







        '''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	)
          model.to(__UpperCamelCase	)
          model.eval()
          # make random mask reproducible
          torch.manual_seed(2	)
          with torch.no_grad():
            _a							=	model(**self._prepare_for_class(__UpperCamelCase     ,     __UpperCamelCase	)	)

          _a							=	outputs[0].cpu().numpy()
          _a							=	0

          with tempfile.TemporaryDirectory() as tmpdirname:
            model.save_pretrained(__UpperCamelCase	)
            _a							=	model_class.from_pretrained(__UpperCamelCase	)
            model.to(__UpperCamelCase	)
            # make random mask reproducible
            torch.manual_seed(2	)
            with torch.no_grad():
              _a							=	model(**self._prepare_for_class(__UpperCamelCase     ,     __UpperCamelCase	)	)

            # Make sure we don't have nans
            _a							=	after_outputs[0].cpu().numpy()
            _a							=	0
            _a							=	np.amax(np.abs(out_a - out_a	)	)
            self.assertLessEqual(__UpperCamelCase     ,     1E-5	)






      @unittest.skip(
          reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."	)
      def 			lowerCamelCase__    (    self	)	->Union[str, Any]:







        '''simple docstring'''


        pass






      @unittest.skip(
          reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."	)
      def 			lowerCamelCase__    (    self	)	->int:







        '''simple docstring'''


        pass






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







        '''simple docstring'''


        pass






      @unittest.skip(reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load"	)
      def 			lowerCamelCase__    (    self	)	->Optional[Any]:







        '''simple docstring'''


        pass






      @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."	)
      def 			lowerCamelCase__    (    self	)	->Tuple:







        '''simple docstring'''


        pass






      @slow
      def 			lowerCamelCase__    (    self	)	->Union[str, Any]:







        '''simple docstring'''


        for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          _a							=	ViTMAEModel.from_pretrained(__UpperCamelCase	)
          self.assertIsNotNone(__UpperCamelCase	)







def        lowerCAmelCase (      )       ->							Optional[Any]:







  '''simple docstring'''







  _a							=	Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"					)
  return image




@require_torch
@require_vision
class    lowercase_    (unittest.TestCase     ):






      @cached_property
      def 			lowerCamelCase__    (    self	)	->Any:







        '''simple docstring'''


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






      @slow
      def 			lowerCamelCase__    (    self	)	->str:







        '''simple docstring'''


        np.random.seed(2	)

        _a							=	ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base"	).to(__UpperCamelCase	)

        _a							=	self.default_image_processor
        _a							=	prepare_img()
        _a							=	image_processor(images=__UpperCamelCase     ,     return_tensors="pt"	).to(__UpperCamelCase	)

        # 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
        with torch.no_grad():
          _a							=	model(**__UpperCamelCase     ,     noise=torch.from_numpy(__UpperCamelCase	).to(device=__UpperCamelCase	)	)

        # verify the logits
        _a							=	torch.Size((1, 1_9_6, 7_6_8)	)
        self.assertEqual(outputs.logits.shape     ,     __UpperCamelCase	)

        _a							=	torch.tensor(
            [[-0.0_548, -1.7_023, -0.9_325], [0.3_721, -0.5_670, -0.2_233], [0.8_235, -1.3_878, -0.3_524]]	)

        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]     ,     expected_slice.to(__UpperCamelCase	)     ,     atol=1E-4	)	)









"""simple docstring"""




from __future__ import annotations

from collections.abc import Callable
from typing import Generic, TypeVar

UpperCamelCase							       =    TypeVar("""T""")
UpperCamelCase							       =    TypeVar("""U""")




class    lowercase_    (Generic[T, U]     ):






      def __init__(    self     ,     a_     ,     a_	)	->Any:







        '''simple docstring'''


        _a							=	key
        _a							=	val
        _a							=	None
        _a							=	None






      def __repr__(    self	)	->str:







        '''simple docstring'''


        return (
            f'''Node: key: {self.key}, val: {self.val}, '''
            f'''has next: {bool(self.next	)}, has prev: {bool(self.prev	)}'''
        )







class    lowercase_    (Generic[T, U]     ):






      def __init__(    self	)	->None:







        '''simple docstring'''


        _a							=	DoubleLinkedListNode(a_     ,     a_	)
        _a							=	DoubleLinkedListNode(a_     ,     a_	)
        _a			,	_a							=	self.rear, self.head






      def __repr__(    self	)	->str:







        '''simple docstring'''


        _a							=	["DoubleLinkedList"]
        _a							=	self.head
        while node.next is not None:
          rep.append(str(a_	)	)
          _a							=	node.next
        rep.append(str(self.rear	)	)
        return ",\n    ".join(a_	)






      def 			lowerCamelCase__    (    self     ,     a_	)	->None:







        '''simple docstring'''


        _a							=	self.rear.prev

        # All nodes other than self.head are guaranteed to have non-None previous
        assert previous is not None

        _a							=	node
        _a							=	previous
        _a							=	node
        _a							=	self.rear






      def 			lowerCamelCase__    (    self     ,     a_	)	->DoubleLinkedListNode[T, U] | None:







        '''simple docstring'''


        if node.prev is None or node.next is None:
          return None

        _a							=	node.next
        _a							=	node.prev
        _a							=	None
        _a							=	None
        return node







class    lowercase_    (Generic[T, U]     ):
      A__	:		dict[Callable[[T], U], LRUCache[T, U]]       			=				{}






      def __init__(    self     ,     a_	)	->int:







        '''simple docstring'''


        _a							=	DoubleLinkedList()
        _a							=	capacity
        _a							=	0
        _a							=	0
        _a							=	0
        _a							=	{}






      def __repr__(    self	)	->str:







        '''simple docstring'''


        return (
            f'''CacheInfo(hits={self.hits}, misses={self.miss}, '''
            f'''capacity={self.capacity}, current size={self.num_keys})'''
        )






      def __contains__(    self     ,     a_	)	->bool:







        '''simple docstring'''


        return key in self.cache






      def 			lowerCamelCase__    (    self     ,     a_	)	->U | None:







        '''simple docstring'''


        if key in self.cache:
          self.hits += 1
          _a							=	self.cache[key]
          _a							=	self.list.remove(self.cache[key]	)
          assert node == value_node

          # node is guaranteed not None because it is in self.cache
          assert node is not None
          self.list.add(a_	)
          return node.val
        self.miss += 1
        return None






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_	)	->None:







        '''simple docstring'''


        if key not in self.cache:
          if self.num_keys >= self.capacity:
            # delete first node (oldest) when over capacity
            _a							=	self.list.head.next

            # guaranteed to have a non-None first node when num_keys > 0
            # explain to type checker via assertions
            assert first_node is not None
            assert first_node.key is not None
            assert (
                self.list.remove(a_	) is not None
            )  # node guaranteed to be in list assert node.key is not None

            del self.cache[first_node.key]
            self.num_keys -= 1
          _a							=	DoubleLinkedListNode(a_     ,     a_	)
          self.list.add(self.cache[key]	)
          self.num_keys += 1

        else:
          # bump node to the end of the list, update value
          _a							=	self.list.remove(self.cache[key]	)
          assert node is not None  # node guaranteed to be in list
          _a							=	value
          self.list.add(a_	)






      @classmethod
      def 			lowerCamelCase__    (    cls     ,     a_ = 1_2_8	)	->Callable[[Callable[[T], U]], Callable[..., U]]:







        '''simple docstring'''


        def cache_decorator_inner(a_	) -> Callable[..., U]:
          def cache_decorator_wrapper(*a_	) -> U:
            if func not in cls.decorator_function_to_instance_map:
              _a							=	LRUCache(a_	)

            _a							=	cls.decorator_function_to_instance_map[func].get(args[0]	)
            if result is None:
              _a							=	func(*a_	)
              cls.decorator_function_to_instance_map[func].put(args[0]     ,     a_	)
            return result

          def cache_info() -> LRUCache[T, U]:
            return cls.decorator_function_to_instance_map[func]

          setattr(a_     ,     "cache_info"     ,     a_	)  # noqa: B010

          return cache_decorator_wrapper

        return cache_decorator_inner


if __name__ == "__main__":
     import doctest

     doctest.testmod()


