"""simple docstring"""







import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    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 ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
       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  UpperCamelCase	:
   def __init__(    self  ,__UpperCamelCase  ,__UpperCamelCase=13  ,__UpperCamelCase=30  ,__UpperCamelCase=2  ,__UpperCamelCase=3  ,__UpperCamelCase=True  ,__UpperCamelCase=True  ,__UpperCamelCase=32  ,__UpperCamelCase=5  ,__UpperCamelCase=4  ,__UpperCamelCase=37  ,__UpperCamelCase="gelu"  ,__UpperCamelCase=0.1  ,__UpperCamelCase=0.1  ,__UpperCamelCase=10  ,__UpperCamelCase=0.02  ,__UpperCamelCase=None  ,__UpperCamelCase=2  ,)			->   Any:


        '''simple docstring'''


        lowercase_			:       List[Any]									=			parent
        lowercase_			:       Union[str, Any]									=			batch_size
        lowercase_			:       Any									=			image_size
        lowercase_			:       str									=			patch_size
        lowercase_			:       Optional[Any]									=			num_channels
        lowercase_			:       Optional[Any]									=			is_training
        lowercase_			:       str									=			use_labels
        lowercase_			:       List[Any]									=			hidden_size
        lowercase_			:       List[Any]									=			num_hidden_layers
        lowercase_			:       Any									=			num_attention_heads
        lowercase_			:       Tuple									=			intermediate_size
        lowercase_			:       Optional[int]									=			hidden_act
        lowercase_			:       Union[str, Any]									=			hidden_dropout_prob
        lowercase_			:       int									=			attention_probs_dropout_prob
        lowercase_			:       str									=			type_sequence_label_size
        lowercase_			:       List[Any]									=			initializer_range
        lowercase_			:       int									=			scope
        lowercase_			:       Dict									=			encoder_stride

        # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
        lowercase_			:       List[Any]									=			(image_size // patch_size) ** 2
        lowercase_			:       Union[str, Any]									=			num_patches + 1
   def     _UpperCAmelCase     (    self    )			->   List[str]:


        '''simple docstring'''


        lowercase_			:       Optional[Any]									=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

        lowercase_			:       List[str]									=			None
        if self.use_labels:
             lowercase_			:       Any									=			ids_tensor([self.batch_size]  ,self.type_sequence_label_size    )

        lowercase_			:       Dict									=			self.get_config()

        return config, pixel_values, labels
   def     _UpperCAmelCase     (    self    )			->   List[Any]:


        '''simple docstring'''


        return ViTConfig(
            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  ,encoder_stride=self.encoder_stride  ,)
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )			->   List[str]:


        '''simple docstring'''


        lowercase_			:       Optional[Any]									=			ViTModel(config=__UpperCamelCase    )
        model.to(__UpperCamelCase    )
        model.eval()
        lowercase_			:       Tuple									=			model(__UpperCamelCase    )
        self.parent.assertEqual(result.last_hidden_state.shape  ,(self.batch_size, self.seq_length, self.hidden_size)    )
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )			->   Union[str, Any]:


        '''simple docstring'''


        lowercase_			:       Any									=			ViTForMaskedImageModeling(config=__UpperCamelCase    )
        model.to(__UpperCamelCase    )
        model.eval()
        lowercase_			:       Dict									=			model(__UpperCamelCase    )
        self.parent.assertEqual(
            result.reconstruction.shape  ,(self.batch_size, self.num_channels, self.image_size, self.image_size)    )

        # test greyscale images
        lowercase_			:       List[str]									=			1
        lowercase_			:       str									=			ViTForMaskedImageModeling(__UpperCamelCase    )
        model.to(__UpperCamelCase    )
        model.eval()

        lowercase_			:       Any									=			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
        lowercase_			:       Dict									=			model(__UpperCamelCase    )
        self.parent.assertEqual(result.reconstruction.shape  ,(self.batch_size, 1, self.image_size, self.image_size)    )
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )			->   Optional[int]:


        '''simple docstring'''


        lowercase_			:       str									=			self.type_sequence_label_size
        lowercase_			:       int									=			ViTForImageClassification(__UpperCamelCase    )
        model.to(__UpperCamelCase    )
        model.eval()
        lowercase_			:       Union[str, Any]									=			model(__UpperCamelCase  ,labels=__UpperCamelCase    )
        self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.type_sequence_label_size)    )

        # test greyscale images
        lowercase_			:       List[str]									=			1
        lowercase_			:       int									=			ViTForImageClassification(__UpperCamelCase    )
        model.to(__UpperCamelCase    )
        model.eval()

        lowercase_			:       Tuple									=			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
        lowercase_			:       Union[str, Any]									=			model(__UpperCamelCase    )
        self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.type_sequence_label_size)    )







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


        '''simple docstring'''


        lowercase_			:       Optional[int]									=			self.prepare_config_and_inputs()
        (
            lowercase_						
        )			:       List[Any]									=			config_and_inputs
        lowercase_			:       Union[str, Any]									=			{'pixel_values': pixel_values}
        return config, inputs_dict







@require_torch
class  UpperCamelCase	(    lowercase_		,   lowercase_		,   unittest.TestCase						):
   lowercase							=						(
       (
           ViTModel,
           ViTForImageClassification,
           ViTForMaskedImageModeling,
       )
       if is_torch_available()
       else ()
   )
   lowercase							=						(
       {'feature-extraction': ViTModel, 'image-classification': ViTForImageClassification}
       if is_torch_available()
       else {}
   )
   lowercase							=						True

   lowercase							=						False
   lowercase							=						False
   lowercase							=						False
   def     _UpperCAmelCase     (    self    )			->   Optional[Any]:


        '''simple docstring'''


        lowercase_			:       str									=			ViTModelTester(self    )
        lowercase_			:       List[Any]									=			ConfigTester(self  ,config_class=__UpperCamelCase  ,has_text_modality=__UpperCamelCase  ,hidden_size=37    )
   def     _UpperCAmelCase     (    self    )			->   List[Any]:


        '''simple docstring'''


        self.config_tester.run_common_tests()
   @unittest.skip(reason='ViT does not use inputs_embeds'    )
   def     _UpperCAmelCase     (    self    )			->   Optional[Any]:


        '''simple docstring'''


        pass
   def     _UpperCAmelCase     (    self    )			->   Optional[int]:


        '''simple docstring'''


        lowercase_			:       Any									=			self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
             lowercase_			:       Tuple									=			model_class(__UpperCamelCase    )
             self.assertIsInstance(model.get_input_embeddings()  ,(nn.Module)    )
             lowercase_			:       Any									=			model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(__UpperCamelCase  ,nn.Linear    )    )
   def     _UpperCAmelCase     (    self    )			->   Dict:


        '''simple docstring'''


        lowercase_			:       Union[str, Any]									=			self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
             lowercase_			:       Union[str, Any]									=			model_class(__UpperCamelCase    )
             lowercase_			:       str									=			inspect.signature(model.forward    )
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             lowercase_			:       List[str]									=			[*signature.parameters.keys()]

             lowercase_			:       List[str]									=			['pixel_values']
             self.assertListEqual(arg_names[:1]  ,__UpperCamelCase    )
   def     _UpperCAmelCase     (    self    )			->   List[str]:


        '''simple docstring'''


        lowercase_			:       Any									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__UpperCamelCase    )
   def     _UpperCAmelCase     (    self    )			->   Optional[int]:


        '''simple docstring'''


        lowercase_			:       Union[str, Any]									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_image_modeling(*__UpperCamelCase    )
   def     _UpperCAmelCase     (    self    )			->   Optional[Any]:


        '''simple docstring'''


        lowercase_			:       Optional[int]									=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*__UpperCamelCase    )







   @slow
   def     _UpperCAmelCase     (    self    )			->   Tuple:


        '''simple docstring'''


        for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowercase_			:       Optional[Any]									=			ViTModel.from_pretrained(__UpperCamelCase    )
             self.assertIsNotNone(__UpperCamelCase    )






def     lowercase__( ):
     lowercase_			:       Tuple									=			Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
     return image







@require_torch
@require_vision
class  UpperCamelCase	(    unittest.TestCase						):
   @cached_property
   def     _UpperCAmelCase     (    self    )			->   List[str]:


        '''simple docstring'''


        return ViTImageProcessor.from_pretrained('google/vit-base-patch16-224'    ) if is_vision_available() else None
   @slow
   def     _UpperCAmelCase     (    self    )			->   str:


        '''simple docstring'''


        lowercase_			:       List[Any]									=			ViTForImageClassification.from_pretrained('google/vit-base-patch16-224'    ).to(__UpperCamelCase    )

        lowercase_			:       int									=			self.default_image_processor
        lowercase_			:       Optional[int]									=			prepare_img()
        lowercase_			:       Any									=			image_processor(images=__UpperCamelCase  ,return_tensors='pt'    ).to(__UpperCamelCase    )

        # forward pass
        with torch.no_grad():
             lowercase_			:       Tuple									=			model(**__UpperCamelCase    )

        # verify the logits
        lowercase_			:       Optional[int]									=			torch.Size((1, 1000)    )
        self.assertEqual(outputs.logits.shape  ,__UpperCamelCase    )

        lowercase_			:       List[str]									=			torch.tensor([-0.2744, 0.8215, -0.0836]    ).to(__UpperCamelCase    )

        self.assertTrue(torch.allclose(outputs.logits[0, :3]  ,__UpperCamelCase  ,atol=1e-4    )    )
   @slow
   def     _UpperCAmelCase     (    self    )			->   int:


        '''simple docstring'''


        lowercase_			:       Optional[int]									=			ViTModel.from_pretrained('facebook/dino-vits8'    ).to(__UpperCamelCase    )

        lowercase_			:       Optional[int]									=			ViTImageProcessor.from_pretrained('facebook/dino-vits8'  ,size=480    )
        lowercase_			:       Union[str, Any]									=			prepare_img()
        lowercase_			:       Optional[Any]									=			image_processor(images=__UpperCamelCase  ,return_tensors='pt'    )
        lowercase_			:       Tuple									=			inputs.pixel_values.to(__UpperCamelCase    )

        # forward pass
        with torch.no_grad():
             lowercase_			:       List[str]									=			model(__UpperCamelCase  ,interpolate_pos_encoding=__UpperCamelCase    )

        # verify the logits
        lowercase_			:       Optional[Any]									=			torch.Size((1, 3601, 384)    )
        self.assertEqual(outputs.last_hidden_state.shape  ,__UpperCamelCase    )

        lowercase_			:       Tuple									=			torch.tensor(
            [[4.2340, 4.3906, -6.6692], [4.5463, 1.8928, -6.7257], [4.4429, 0.8496, -5.8585]]    ).to(__UpperCamelCase    )

        self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]  ,__UpperCamelCase  ,atol=1e-4    )    )







   @slow
   @require_accelerate
   @require_torch_gpu
   def     _UpperCAmelCase     (    self    )			->   Dict:


        '''simple docstring'''


        lowercase_			:       Any									=			ViTModel.from_pretrained('facebook/dino-vits8'  ,torch_dtype=torch.floataa  ,device_map='auto'    )
        lowercase_			:       Optional[int]									=			self.default_image_processor

        lowercase_			:       int									=			prepare_img()
        lowercase_			:       List[str]									=			image_processor(images=__UpperCamelCase  ,return_tensors='pt'    )
        lowercase_			:       Union[str, Any]									=			inputs.pixel_values.to(__UpperCamelCase    )

        # forward pass to make sure inference works in fp16
        with torch.no_grad():
             lowercase_			:       Any									=			model(__UpperCamelCase    )


"""simple docstring"""







import os
import unittest
from tempfile import TemporaryDirectory

import torch
import torch.nn as nn

from accelerate.utils import (
    OffloadedWeightsLoader,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
)







class  UpperCamelCase	(    nn.Module						):
   def __init__(    self    )			->   Any:


        '''simple docstring'''


        super().__init__()
        lowercase_			:       Dict									=			nn.Linear(3  ,4    )
        lowercase_			:       Union[str, Any]									=			nn.BatchNormad(4    )
        lowercase_			:       Union[str, Any]									=			nn.Linear(4  ,5    )







   def     _UpperCAmelCase     (    self  ,__UpperCamelCase    )			->   int:


        '''simple docstring'''


        return self.lineara(self.batchnorm(self.lineara(__UpperCamelCase    )    )    )







class  UpperCamelCase	(    unittest.TestCase						):
   def     _UpperCAmelCase     (    self    )			->   Dict:


        '''simple docstring'''


        lowercase_			:       int									=			ModelForTest()
        with TemporaryDirectory() as tmp_dir:
             offload_state_dict(__UpperCamelCase  ,model.state_dict()    )
             lowercase_			:       str									=			os.path.join(__UpperCamelCase  ,'index.json'    )
             self.assertTrue(os.path.isfile(__UpperCamelCase    )    )
             # TODO: add tests on what is inside the index

             for key in ["linear1.weight", "linear1.bias", "linear2.weight", "linear2.bias"]:
                  lowercase_			:       int									=			os.path.join(__UpperCamelCase  ,f'''{key}.dat'''    )
                  self.assertTrue(os.path.isfile(__UpperCamelCase    )    )
                  # TODO: add tests on the fact weights are properly loaded
   def     _UpperCAmelCase     (    self    )			->   Tuple:


        '''simple docstring'''


        lowercase_			:       List[str]									=			[torch.floataa, torch.floataa, torch.bfloataa]

        for dtype in dtypes:
             lowercase_			:       Tuple									=			torch.randn(2  ,3  ,dtype=__UpperCamelCase    )
             with TemporaryDirectory() as tmp_dir:
                  lowercase_			:       List[str]									=			offload_weight(__UpperCamelCase  ,'weight'  ,__UpperCamelCase  ,{}    )
                  lowercase_			:       str									=			os.path.join(__UpperCamelCase  ,'weight.dat'    )
                  self.assertTrue(os.path.isfile(__UpperCamelCase    )    )
                  self.assertDictEqual(__UpperCamelCase  ,{'weight': {'shape': [2, 3], 'dtype': str(__UpperCamelCase    ).split('.'    )[1]}}    )

                  lowercase_			:       Optional[Any]									=			load_offloaded_weight(__UpperCamelCase  ,index['weight']    )
                  self.assertTrue(torch.equal(__UpperCamelCase  ,__UpperCamelCase    )    )
   def     _UpperCAmelCase     (    self    )			->   List[Any]:


        '''simple docstring'''


        lowercase_			:       List[str]									=			ModelForTest()
        lowercase_			:       Optional[int]									=			model.state_dict()
        lowercase_			:       int									=			{k: v for k, v in state_dict.items() if 'linear2' not in k}
        lowercase_			:       Optional[Any]									=			{k: v for k, v in state_dict.items() if 'linear2' in k}

        with TemporaryDirectory() as tmp_dir:
             offload_state_dict(__UpperCamelCase  ,__UpperCamelCase    )
             lowercase_			:       Tuple									=			OffloadedWeightsLoader(state_dict=__UpperCamelCase  ,save_folder=__UpperCamelCase    )

             # Every key is there with the right value
             self.assertEqual(sorted(__UpperCamelCase    )  ,sorted(state_dict.keys()    )    )
             for key, param in state_dict.items():
                  self.assertTrue(torch.allclose(__UpperCamelCase  ,weight_map[key]    )    )

        lowercase_			:       int									=			{k: v for k, v in state_dict.items() if 'weight' in k}
        lowercase_			:       int									=			{k: v for k, v in state_dict.items() if 'weight' not in k}

        with TemporaryDirectory() as tmp_dir:
             offload_state_dict(__UpperCamelCase  ,__UpperCamelCase    )
             lowercase_			:       int									=			OffloadedWeightsLoader(state_dict=__UpperCamelCase  ,save_folder=__UpperCamelCase    )

             # Every key is there with the right value
             self.assertEqual(sorted(__UpperCamelCase    )  ,sorted(state_dict.keys()    )    )
             for key, param in state_dict.items():
                  self.assertTrue(torch.allclose(__UpperCamelCase  ,weight_map[key]    )    )

        with TemporaryDirectory() as tmp_dir:
             offload_state_dict(__UpperCamelCase  ,__UpperCamelCase    )
             # Duplicates are removed
             lowercase_			:       Tuple									=			OffloadedWeightsLoader(state_dict=__UpperCamelCase  ,save_folder=__UpperCamelCase    )

             # Every key is there with the right value
             self.assertEqual(sorted(__UpperCamelCase    )  ,sorted(state_dict.keys()    )    )
             for key, param in state_dict.items():
                  self.assertTrue(torch.allclose(__UpperCamelCase  ,weight_map[key]    )    )







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


        '''simple docstring'''


        lowercase_			:       Optional[int]									=			{'a.1': 0, 'a.10': 1, 'a.2': 2}
        lowercase_			:       int									=			extract_submodules_state_dict(__UpperCamelCase  ,['a.1', 'a.2']    )
        self.assertDictEqual(__UpperCamelCase  ,{'a.1': 0, 'a.2': 2}    )

        lowercase_			:       Union[str, Any]									=			{'a.1.a': 0, 'a.10.a': 1, 'a.2.a': 2}
        lowercase_			:       Optional[int]									=			extract_submodules_state_dict(__UpperCamelCase  ,['a.1', 'a.2']    )
        self.assertDictEqual(__UpperCamelCase  ,{'a.1.a': 0, 'a.2.a': 2}    )

