





'''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       __SCREAMING_SNAKE_CASE	:



     '''simple docstring'''






     def __init__(					self						,   lowerCamelCase						,   lowerCamelCase=13						,   lowerCamelCase=30						,   lowerCamelCase=2						,   lowerCamelCase=3						,   lowerCamelCase=True						,   lowerCamelCase=True						,   lowerCamelCase=32						,   lowerCamelCase=5						,   lowerCamelCase=4						,   lowerCamelCase=37						,   lowerCamelCase="gelu"						,   lowerCamelCase=0.1						,   lowerCamelCase=0.1						,   lowerCamelCase=10						,   lowerCamelCase=0.02						,   lowerCamelCase=None						,   lowerCamelCase=2						,   ):
       _snake_case     		=				parent
       _snake_case     		=				batch_size
       _snake_case     		=				image_size
       _snake_case     		=				patch_size
       _snake_case     		=				num_channels
       _snake_case     		=				is_training
       _snake_case     		=				use_labels
       _snake_case     		=				hidden_size
       _snake_case     		=				num_hidden_layers
       _snake_case     		=				num_attention_heads
       _snake_case     		=				intermediate_size
       _snake_case     		=				hidden_act
       _snake_case     		=				hidden_dropout_prob
       _snake_case     		=				attention_probs_dropout_prob
       _snake_case     		=				type_sequence_label_size
       _snake_case     		=				initializer_range
       _snake_case     		=				scope
       _snake_case     		=				encoder_stride

       # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
       _snake_case     		=				(image_size // patch_size) ** 2
       _snake_case     		=				num_patches + 1

     def 	UpperCamelCase(					self	):
       _snake_case     		=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

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

       _snake_case     		=				self.get_config()

       return config, pixel_values, labels

     def 	UpperCamelCase(					self	):
       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=lowerCamelCase						,   initializer_range=self.initializer_range						,   encoder_stride=self.encoder_stride						,   )

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase	):
       _snake_case     		=				ViTModel(config=lowerCamelCase	)
       model.to(lowerCamelCase	)
       model.eval()
       _snake_case     		=				model(lowerCamelCase	)
       self.parent.assertEqual(result.last_hidden_state.shape						,   (self.batch_size, self.seq_length, self.hidden_size)	)

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase	):
       _snake_case     		=				ViTForMaskedImageModeling(config=lowerCamelCase	)
       model.to(lowerCamelCase	)
       model.eval()
       _snake_case     		=				model(lowerCamelCase	)
       self.parent.assertEqual(
           result.reconstruction.shape						,   (self.batch_size, self.num_channels, self.image_size, self.image_size)	)

       # test greyscale images
       _snake_case     		=				1
       _snake_case     		=				ViTForMaskedImageModeling(lowerCamelCase	)
       model.to(lowerCamelCase	)
       model.eval()

       _snake_case     		=				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
       _snake_case     		=				model(lowerCamelCase	)
       self.parent.assertEqual(result.reconstruction.shape						,   (self.batch_size, 1, self.image_size, self.image_size)	)

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase	):
       _snake_case     		=				self.type_sequence_label_size
       _snake_case     		=				ViTForImageClassification(lowerCamelCase	)
       model.to(lowerCamelCase	)
       model.eval()
       _snake_case     		=				model(lowerCamelCase						,   labels=lowerCamelCase	)
       self.parent.assertEqual(result.logits.shape						,   (self.batch_size, self.type_sequence_label_size)	)

       # test greyscale images
       _snake_case     		=				1
       _snake_case     		=				ViTForImageClassification(lowerCamelCase	)
       model.to(lowerCamelCase	)
       model.eval()

       _snake_case     		=				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
       _snake_case     		=				model(lowerCamelCase	)
       self.parent.assertEqual(result.logits.shape						,   (self.batch_size, self.type_sequence_label_size)	)






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







@require_torch
class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase					,    __UpperCamelCase					,    unittest.TestCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						int      =   (
         (
             ViTModel,
             ViTForImageClassification,
             ViTForMaskedImageModeling,
         )
         if is_torch_available()
         else ()
     )
     UpperCAmelCase__       :						Any      =   (
         {'''feature-extraction''': ViTModel, '''image-classification''': ViTForImageClassification}
         if is_torch_available()
         else {}
     )
     UpperCAmelCase__       :						int      =   True

     UpperCAmelCase__       :						Optional[int]      =   False
     UpperCAmelCase__       :						List[str]      =   False
     UpperCAmelCase__       :						Tuple      =   False

     def 	UpperCamelCase(					self	):
       _snake_case     		=				ViTModelTester(self	)
       _snake_case     		=				ConfigTester(self						,   config_class=lowerCamelCase						,   has_text_modality=lowerCamelCase						,   hidden_size=37	)

     def 	UpperCamelCase(					self	):
       self.config_tester.run_common_tests()

     @unittest.skip(reason="ViT does not use inputs_embeds"	)
     def 	UpperCamelCase(					self	):
       pass

     def 	UpperCamelCase(					self	):
       _snake_case		,		_snake_case     		=				self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
         _snake_case     		=				model_class(lowerCamelCase	)
         self.assertIsInstance(model.get_input_embeddings()						,   (nn.Module)	)
         _snake_case     		=				model.get_output_embeddings()
         self.assertTrue(x is None or isinstance(lowerCamelCase						,   nn.Linear	)	)

     def 	UpperCamelCase(					self	):
       _snake_case		,		_snake_case     		=				self.model_tester.prepare_config_and_inputs_for_common()

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

         _snake_case     		=				["pixel_values"]
         self.assertListEqual(arg_names[:1]						,   lowerCamelCase	)

     def 	UpperCamelCase(					self	):
       _snake_case     		=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*lowerCamelCase	)

     def 	UpperCamelCase(					self	):
       _snake_case     		=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_image_modeling(*lowerCamelCase	)

     def 	UpperCamelCase(					self	):
       _snake_case     		=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*lowerCamelCase	)






     @slow
     def 	UpperCamelCase(					self	):
       for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         _snake_case     		=				ViTModel.from_pretrained(lowerCamelCase	)
         self.assertIsNotNone(lowerCamelCase	)







def 				snake_case_      (		):







  '''simple docstring'''



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




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



     '''simple docstring'''






     @cached_property
     def 	UpperCamelCase(					self	):
       return ViTImageProcessor.from_pretrained("google/vit-base-patch16-224"	) if is_vision_available() else None

     @slow
     def 	UpperCamelCase(					self	):
       _snake_case     		=				ViTForImageClassification.from_pretrained("google/vit-base-patch16-224"	).to(lowerCamelCase	)

       _snake_case     		=				self.default_image_processor
       _snake_case     		=				prepare_img()
       _snake_case     		=				image_processor(images=lowerCamelCase						,   return_tensors="pt"	).to(lowerCamelCase	)

       # forward pass
       with torch.no_grad():
         _snake_case     		=				model(**lowerCamelCase	)

       # verify the logits
       _snake_case     		=				torch.Size((1, 1_000)	)
       self.assertEqual(outputs.logits.shape						,   lowerCamelCase	)

       _snake_case     		=				torch.tensor([-0.2744, 0.8215, -0.0836]	).to(lowerCamelCase	)

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

     @slow
     def 	UpperCamelCase(					self	):
       # ViT models have an `interpolate_pos_encoding` argument in their forward method,
       # allowing to interpolate the pre-trained position embeddings in order to use
       # the model on higher resolutions. The DINO model by Facebook AI leverages this
       # to visualize self-attention on higher resolution images.
       _snake_case     		=				ViTModel.from_pretrained("facebook/dino-vits8"	).to(lowerCamelCase	)

       _snake_case     		=				ViTImageProcessor.from_pretrained("facebook/dino-vits8"						,   size=480	)
       _snake_case     		=				prepare_img()
       _snake_case     		=				image_processor(images=lowerCamelCase						,   return_tensors="pt"	)
       _snake_case     		=				inputs.pixel_values.to(lowerCamelCase	)

       # forward pass
       with torch.no_grad():
         _snake_case     		=				model(lowerCamelCase						,   interpolate_pos_encoding=lowerCamelCase	)

       # verify the logits
       _snake_case     		=				torch.Size((1, 3_601, 384)	)
       self.assertEqual(outputs.last_hidden_state.shape						,   lowerCamelCase	)

       _snake_case     		=				torch.tensor(
           [[4.2340, 4.3906, -6.6692], [4.5463, 1.8928, -6.7257], [4.4429, 0.8496, -5.8585]]	).to(lowerCamelCase	)

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






     @slow
     @require_accelerate
     @require_torch_gpu
     def 	UpperCamelCase(					self	):
       _snake_case     		=				ViTModel.from_pretrained("facebook/dino-vits8"						,   torch_dtype=torch.floataa						,   device_map="auto"	)
       _snake_case     		=				self.default_image_processor

       _snake_case     		=				prepare_img()
       _snake_case     		=				image_processor(images=lowerCamelCase						,   return_tensors="pt"	)
       _snake_case     		=				inputs.pixel_values.to(lowerCamelCase	)

       # forward pass to make sure inference works in fp16
       with torch.no_grad():
         _snake_case     		=				model(lowerCamelCase	)









'''simple docstring'''



import os
import tempfile
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
 import torch
 from torch import nn

 from transformers import (
     Adafactor,
     AdamW,
     get_constant_schedule,
     get_constant_schedule_with_warmup,
     get_cosine_schedule_with_warmup,
     get_cosine_with_hard_restarts_schedule_with_warmup,
     get_inverse_sqrt_schedule,
     get_linear_schedule_with_warmup,
     get_polynomial_decay_schedule_with_warmup,
 )




def 				snake_case_      (		SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__=10       ):







  '''simple docstring'''



  _snake_case     		=				[]
  for _ in range(SCREAMING_SNAKE_CASE__       ):
    lrs.append(scheduler.get_lr()[0]       )
    scheduler.step()
  return lrs




def 				snake_case_      (		SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__=10       ):







  '''simple docstring'''



  _snake_case     		=				[]
  for step in range(SCREAMING_SNAKE_CASE__       ):
    lrs.append(scheduler.get_lr()[0]       )
    scheduler.step()
    if step == num_steps // 2:
      with tempfile.TemporaryDirectory() as tmpdirname:
        _snake_case     		=				os.path.join(SCREAMING_SNAKE_CASE__					,	"schedule.bin"       )
        torch.save(scheduler.state_dict()					,	SCREAMING_SNAKE_CASE__       )

        _snake_case     		=				torch.load(SCREAMING_SNAKE_CASE__       )
        scheduler.load_state_dict(SCREAMING_SNAKE_CASE__       )
  return lrs


@require_torch
class       __SCREAMING_SNAKE_CASE	(						unittest.TestCase			):



     '''simple docstring'''






     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase	):
       self.assertEqual(len(lowerCamelCase	)						,   len(lowerCamelCase	)	)
       for a, b in zip(lowerCamelCase						,   lowerCamelCase	):
         self.assertAlmostEqual(lowerCamelCase						,   lowerCamelCase						,   delta=lowerCamelCase	)

     def 	UpperCamelCase(					self	):
       _snake_case     		=				torch.tensor([0.1, -0.2, -0.1]						,   requires_grad=lowerCamelCase	)
       _snake_case     		=				torch.tensor([0.4, 0.2, -0.5]	)
       _snake_case     		=				nn.MSELoss()
       # No warmup, constant schedule, no gradient clipping
       _snake_case     		=				AdamW(params=[w]						,   lr=2e-1						,   weight_decay=0.0	)
       for _ in range(100	):
         _snake_case     		=				criterion(lowerCamelCase						,   lowerCamelCase	)
         loss.backward()
         optimizer.step()
         w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
         w.grad.zero_()
       self.assertListAlmostEqual(w.tolist()						,   [0.4, 0.2, -0.5]						,   tol=1e-2	)






     def 	UpperCamelCase(					self	):
       _snake_case     		=				torch.tensor([0.1, -0.2, -0.1]						,   requires_grad=lowerCamelCase	)
       _snake_case     		=				torch.tensor([0.4, 0.2, -0.5]	)
       _snake_case     		=				nn.MSELoss()
       # No warmup, constant schedule, no gradient clipping
       _snake_case     		=				Adafactor(
           params=[w]						,   lr=1e-2						,   eps=(1e-30, 1e-3)						,   clip_threshold=1.0						,   decay_rate=-0.8						,   betaa=lowerCamelCase						,   weight_decay=0.0						,   relative_step=lowerCamelCase						,   scale_parameter=lowerCamelCase						,   warmup_init=lowerCamelCase						,   )
       for _ in range(1_000	):
         _snake_case     		=				criterion(lowerCamelCase						,   lowerCamelCase	)
         loss.backward()
         optimizer.step()
         w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
         w.grad.zero_()
       self.assertListAlmostEqual(w.tolist()						,   [0.4, 0.2, -0.5]						,   tol=1e-2	)







@require_torch
class       __SCREAMING_SNAKE_CASE	(						unittest.TestCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						Any      =   nn.Linear(50					,    50			) if is_torch_available() else None
     UpperCAmelCase__       :						str      =   AdamW(m.parameters()					,    lr=10.0			) if is_torch_available() else None
     UpperCAmelCase__       :						Optional[Any]      =   10

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase=None	):
       self.assertEqual(len(lowerCamelCase	)						,   len(lowerCamelCase	)	)
       for a, b in zip(lowerCamelCase						,   lowerCamelCase	):
         self.assertAlmostEqual(lowerCamelCase						,   lowerCamelCase						,   delta=lowerCamelCase						,   msg=lowerCamelCase	)






     def 	UpperCamelCase(					self	):
       _snake_case     		=				{"num_warmup_steps": 2, "num_training_steps": 10}
       # schedulers doct format
       # function: (sched_args_dict, expected_learning_rates)
       _snake_case     		=				{
           get_constant_schedule: ({}, [10.0] * self.num_steps),
           get_constant_schedule_with_warmup: (
               {"num_warmup_steps": 4},
               [0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0],
           ),
           get_linear_schedule_with_warmup: (
               {**common_kwargs},
               [0.0, 5.0, 10.0, 8.75, 7.5, 6.25, 5.0, 3.75, 2.5, 1.25],
           ),
           get_cosine_schedule_with_warmup: (
               {**common_kwargs},
               [0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38],
           ),
           get_cosine_with_hard_restarts_schedule_with_warmup: (
               {**common_kwargs, "num_cycles": 2},
               [0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46],
           ),
           get_polynomial_decay_schedule_with_warmup: (
               {**common_kwargs, "power": 2.0, "lr_end": 1e-7},
               [0.0, 5.0, 10.0, 7.656, 5.625, 3.906, 2.5, 1.406, 0.625, 0.156],
           ),
           get_inverse_sqrt_schedule: (
               {"num_warmup_steps": 2},
               [0.0, 5.0, 10.0, 8.165, 7.071, 6.325, 5.774, 5.345, 5.0, 4.714],
           ),
       }

       for scheduler_func, data in scheds.items():
         _snake_case		,		_snake_case     		=				data

         _snake_case     		=				scheduler_func(self.optimizer						,   **lowerCamelCase	)
         self.assertEqual(len([scheduler.get_lr()[0]]	)						,   1	)
         _snake_case     		=				unwrap_schedule(lowerCamelCase						,   self.num_steps	)
         self.assertListAlmostEqual(
             lowerCamelCase						,   lowerCamelCase						,   tol=1e-2						,   msg=F'''failed for {scheduler_func} in normal scheduler'''						,   )

         _snake_case     		=				scheduler_func(self.optimizer						,   **lowerCamelCase	)
         if scheduler_func.__name__ != "get_constant_schedule":
           LambdaScheduleWrapper.wrap_scheduler(lowerCamelCase	)  # wrap to test picklability of the schedule
         _snake_case     		=				unwrap_and_save_reload_schedule(lowerCamelCase						,   self.num_steps	)
         self.assertListEqual(lowerCamelCase						,   lowerCamelCase						,   msg=F'''failed for {scheduler_func} in save and reload'''	)







class       __SCREAMING_SNAKE_CASE	:



     '''simple docstring'''






     def __init__(					self						,   lowerCamelCase	):
       _snake_case     		=				fn

     def __call__(					self						,   *lowerCamelCase						,   **lowerCamelCase	):
       return self.fn(*lowerCamelCase						,   **lowerCamelCase	)






     @classmethod
     def 	UpperCamelCase(					self						,   lowerCamelCase	):
       _snake_case     		=				list(map(self						,   scheduler.lr_lambdas	)	)


