





'''simple docstring'''

import inspect
import unittest
import warnings
from math import ceil, floor

from transformers import LevitConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

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 transformers import (
       MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
       MODEL_MAPPING,
       LevitForImageClassification,
       LevitForImageClassificationWithTeacher,
       LevitModel,
   )
   from transformers.models.levit.modeling_levit import LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
   from PIL import Image

   from transformers import LevitImageProcessor

class       UpperCAmelCase (						_A  ):




  def 				UpperCAmelCase__   (self      :	List[str]				)					-> List[str]:
      lowercase         = self.config_class(**self.inputs_dict				)
      self.parent.assertTrue(hasattr(A__							,			"hidden_sizes"				)				)
      self.parent.assertTrue(hasattr(A__							,			"num_attention_heads"				)				)




class       UpperCAmelCase :




  def __init__(self      :	List[str]							,			A__      :	Optional[int]							,			A__      :	Tuple=1_3							,			A__      :	List[Any]=6_4							,			A__      :	str=3							,			A__      :	Optional[int]=3							,			A__      :	int=2							,			A__      :	Optional[int]=1							,			A__      :	Optional[Any]=1_6							,			A__      :	Any=[1_2_8, 2_5_6, 3_8_4]							,			A__      :	Optional[Any]=[4, 6, 8]							,			A__      :	Union[str, Any]=[2, 3, 4]							,			A__      :	Any=[1_6, 1_6, 1_6]							,			A__      :	Optional[int]=0							,			A__      :	Optional[int]=[2, 2, 2]							,			A__      :	Union[str, Any]=[2, 2, 2]							,			A__      :	Optional[Any]=0.0_2							,			A__      :	Optional[int]=True							,			A__      :	int=True							,			A__      :	Union[str, Any]=2							,			)					-> List[Any]:
      lowercase         = parent
      lowercase         = batch_size
      lowercase         = image_size
      lowercase         = num_channels
      lowercase         = kernel_size
      lowercase         = stride
      lowercase         = padding
      lowercase         = hidden_sizes
      lowercase         = num_attention_heads
      lowercase         = depths
      lowercase         = key_dim
      lowercase         = drop_path_rate
      lowercase         = patch_size
      lowercase         = attention_ratio
      lowercase         = mlp_ratio
      lowercase         = initializer_range
      lowercase         = [
          ["Subsample", key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
          ["Subsample", key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
      ]
      lowercase         = is_training
      lowercase         = use_labels
      lowercase         = num_labels
      lowercase         = initializer_range




  def 				UpperCAmelCase__   (self      :	Dict				)					-> Optional[int]:
      lowercase         = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

      lowercase         = None
      if self.use_labels:
          lowercase         = ids_tensor([self.batch_size]							,			self.num_labels				)

      lowercase         = self.get_config()
      return config, pixel_values, labels




  def 				UpperCAmelCase__   (self      :	Dict				)					-> List[Any]:
      return LevitConfig(
          image_size=self.image_size							,			num_channels=self.num_channels							,			kernel_size=self.kernel_size							,			stride=self.stride							,			padding=self.padding							,			patch_size=self.patch_size							,			hidden_sizes=self.hidden_sizes							,			num_attention_heads=self.num_attention_heads							,			depths=self.depths							,			key_dim=self.key_dim							,			drop_path_rate=self.drop_path_rate							,			mlp_ratio=self.mlp_ratio							,			attention_ratio=self.attention_ratio							,			initializer_range=self.initializer_range							,			down_ops=self.down_ops							,			)




  def 				UpperCAmelCase__   (self      :	int							,			A__      :	Tuple							,			A__      :	Any							,			A__      :	Optional[Any]				)					-> int:
      lowercase         = LevitModel(config=A__				)
      model.to(A__				)
      model.eval()
      lowercase         = model(A__				)
      lowercase         = (self.image_size, self.image_size)
      lowercase         = image_size[0], image_size[1]
      for _ in range(4				):
          lowercase         = floor(((height + 2 * self.padding - self.kernel_size) / self.stride) + 1				)
          lowercase         = floor(((width + 2 * self.padding - self.kernel_size) / self.stride) + 1				)
      self.parent.assertEqual(
          result.last_hidden_state.shape							,			(self.batch_size, ceil(height / 4				) * ceil(width / 4				), self.hidden_sizes[-1])							,			)




  def 				UpperCAmelCase__   (self      :	Dict							,			A__      :	int							,			A__      :	str							,			A__      :	str				)					-> Dict:
      lowercase         = self.num_labels
      lowercase         = LevitForImageClassification(A__				)
      model.to(A__				)
      model.eval()
      lowercase         = model(A__							,			labels=A__				)
      self.parent.assertEqual(result.logits.shape							,			(self.batch_size, self.num_labels)				)





  def 				UpperCAmelCase__   (self      :	int				)					-> str:
      lowercase         = self.prepare_config_and_inputs()
      lowercase         = config_and_inputs
      lowercase         = {"pixel_values": pixel_values}
      return config, inputs_dict




@require_torch
class       UpperCAmelCase (						_A			,       _A			,       unittest.TestCase  ):
  UpperCAmelCase					:	Dict 			=		(
      (LevitModel, LevitForImageClassification, LevitForImageClassificationWithTeacher)
      if is_torch_available()
      else ()
  )
  UpperCAmelCase					:	Tuple 			=		(
      {
          '''feature-extraction''': LevitModel,
          '''image-classification''': (LevitForImageClassification, LevitForImageClassificationWithTeacher),
      }
      if is_torch_available()
      else {}
  )

  UpperCAmelCase					:	List[str] 			=		False
  UpperCAmelCase					:	str 			=		False
  UpperCAmelCase					:	Tuple 			=		False
  UpperCAmelCase					:	List[Any] 			=		False
  UpperCAmelCase					:	Tuple 			=		False




  def 				UpperCAmelCase__   (self      :	Any				)					-> Any:
      lowercase         = LevitModelTester(self				)
      lowercase         = ConfigTester(self							,			config_class=A__							,			has_text_modality=A__							,			hidden_size=3_7				)




  def 				UpperCAmelCase__   (self      :	List[str]				)					-> Any:
      self.create_and_test_config_common_properties()
      self.config_tester.create_and_test_config_to_json_string()
      self.config_tester.create_and_test_config_to_json_file()
      self.config_tester.create_and_test_config_from_and_save_pretrained()
      self.config_tester.create_and_test_config_with_num_labels()
      self.config_tester.check_config_can_be_init_without_params()
      self.config_tester.check_config_arguments_init()




  def 				UpperCAmelCase__   (self      :	List[str]				)					-> Dict:
      return




  @unittest.skip(reason="Levit does not use inputs_embeds"				)
  def 				UpperCAmelCase__   (self      :	List[Any]				)					-> Any:
      pass




  @unittest.skip(reason="Levit does not support input and output embeddings"				)
  def 				UpperCAmelCase__   (self      :	Optional[Any]				)					-> Dict:
      pass




  @unittest.skip(reason="Levit does not output attentions"				)
  def 				UpperCAmelCase__   (self      :	List[str]				)					-> int:
      pass




  def 				UpperCAmelCase__   (self      :	int				)					-> List[Any]:
      lowercase         = self.model_tester.prepare_config_and_inputs_for_common()

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

          lowercase         = ["pixel_values"]
          self.assertListEqual(arg_names[:1]							,			A__				)




  def 				UpperCAmelCase__   (self      :	Tuple				)					-> Optional[int]:
      def check_hidden_states_output(A__      :	str							,			A__      :	Optional[int]							,			A__      :	Union[str, Any]				):
          lowercase         = model_class(A__				)
          model.to(A__				)
          model.eval()

          with torch.no_grad():
              lowercase         = model(**self._prepare_for_class(A__							,			A__				)				)

          lowercase         = outputs.hidden_states

          lowercase         = len(self.model_tester.depths				) + 1
          self.assertEqual(len(A__				)							,			A__				)

          lowercase         = (self.model_tester.image_size, self.model_tester.image_size)
          lowercase         = image_size[0], image_size[1]
          for _ in range(4				):
              lowercase         = floor(
                  (
                      (height + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                      / self.model_tester.stride
                  )
                  + 1				)
              lowercase         = floor(
                  (
                      (width + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                      / self.model_tester.stride
                  )
                  + 1				)
          # verify the first hidden states (first block)
          self.assertListEqual(
              list(hidden_states[0].shape[-2:]				)							,			[
                  height * width,
                  self.model_tester.hidden_sizes[0],
              ]							,			)

      lowercase         = self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
          lowercase         = True
          check_hidden_states_output(A__							,			A__							,			A__				)

          # check that output_hidden_states also work using config
          del inputs_dict["output_hidden_states"]
          lowercase         = True

          check_hidden_states_output(A__							,			A__							,			A__				)




  @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."				)
  def 				UpperCAmelCase__   (self      :	List[Any]				)					-> Optional[Any]:
      pass




  def 				UpperCAmelCase__   (self      :	int							,			A__      :	List[Any]							,			A__      :	Optional[int]							,			A__      :	List[str]=False				)					-> str:
      lowercase         = super()._prepare_for_class(A__							,			A__							,			return_labels=A__				)

      if return_labels:
          if model_class.__name__ == "LevitForImageClassificationWithTeacher":
              del inputs_dict["labels"]

      return inputs_dict




  def 				UpperCAmelCase__   (self      :	List[str]				)					-> Any:
      lowercase         = self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*A__				)




  def 				UpperCAmelCase__   (self      :	Optional[int]				)					-> Optional[Any]:
      lowercase         = self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*A__				)




  def 				UpperCAmelCase__   (self      :	int				)					-> List[str]:
      if not self.model_tester.is_training:
          return

      lowercase         = self.model_tester.prepare_config_and_inputs_for_common()
      lowercase         = True

      for model_class in self.all_model_classes:
          # LevitForImageClassificationWithTeacher supports inference-only
          if (
              model_class in get_values(A__				)
              or model_class.__name__ == "LevitForImageClassificationWithTeacher"
          ):
              continue
          lowercase         = model_class(A__				)
          model.to(A__				)
          model.train()
          lowercase         = self._prepare_for_class(A__							,			A__							,			return_labels=A__				)
          lowercase         = model(**A__				).loss
          loss.backward()




  def 				UpperCAmelCase__   (self      :	List[str]				)					-> Dict:
      lowercase         = self.model_tester.prepare_config_and_inputs_for_common()
      if not self.model_tester.is_training:
          return

      lowercase         = False
      lowercase         = True

      for model_class in self.all_model_classes:
          if model_class in get_values(A__				) or not model_class.supports_gradient_checkpointing:
              continue
          # LevitForImageClassificationWithTeacher supports inference-only
          if model_class.__name__ == "LevitForImageClassificationWithTeacher":
              continue
          lowercase         = model_class(A__				)
          model.gradient_checkpointing_enable()
          model.to(A__				)
          model.train()
          lowercase         = self._prepare_for_class(A__							,			A__							,			return_labels=A__				)
          lowercase         = model(**A__				).loss
          loss.backward()




  def 				UpperCAmelCase__   (self      :	Optional[Any]				)					-> Optional[Any]:
      lowercase         = self.model_tester.prepare_config_and_inputs_for_common()

      lowercase         = [
          {"title": "multi_label_classification", "num_labels": 2, "dtype": torch.float},
          {"title": "single_label_classification", "num_labels": 1, "dtype": torch.long},
          {"title": "regression", "num_labels": 1, "dtype": torch.float},
      ]

      for model_class in self.all_model_classes:
          if (
              model_class
              not in [
                  *get_values(A__				),
              ]
              or model_class.__name__ == "LevitForImageClassificationWithTeacher"
          ):
              continue

          for problem_type in problem_types:
              with self.subTest(msg=f'Testing {model_class} with {problem_type["title"]}'				):
                  lowercase         = problem_type["title"]
                  lowercase         = problem_type["num_labels"]

                  lowercase         = model_class(A__				)
                  model.to(A__				)
                  model.train()

                  lowercase         = self._prepare_for_class(A__							,			A__							,			return_labels=A__				)

                  if problem_type["num_labels"] > 1:
                      lowercase         = inputs["labels"].unsqueeze(1				).repeat(1							,			problem_type["num_labels"]				)

                  lowercase         = inputs["labels"].to(problem_type["dtype"]				)

                  # This tests that we do not trigger the warning form PyTorch "Using a target size that is different
                  # to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
                  # they have the same size." which is a symptom something in wrong for the regression problem.
                  # See https://github.com/huggingface/transformers/issues/11780
                  with warnings.catch_warnings(record=A__				) as warning_list:
                      lowercase         = model(**A__				).loss
                  for w in warning_list:
                      if "Using a target size that is different to the input size" in str(w.message				):
                          raise ValueError(
                              f'Something is going wrong in the regression problem: intercepted {w.message}'				)

                  loss.backward()





  @slow
  def 				UpperCAmelCase__   (self      :	Union[str, Any]				)					-> Optional[int]:
      for model_name in LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowercase         = LevitModel.from_pretrained(A__				)
          self.assertIsNotNone(A__				)







def 							UpperCAmelCase_    (	):

    """simple docstring"""






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


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




  @cached_property
  def 				UpperCAmelCase__   (self      :	Any				)					-> Tuple:
      return LevitImageProcessor.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]				)





  @slow
  def 				UpperCAmelCase__   (self      :	List[str]				)					-> Any:
      lowercase         = LevitForImageClassificationWithTeacher.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]				).to(
          A__				)

      lowercase         = self.default_image_processor
      lowercase         = prepare_img()
      lowercase         = image_processor(images=A__							,			return_tensors="pt"				).to(A__				)

      # forward pass
      with torch.no_grad():
          lowercase         = model(**A__				)

      # verify the logits
      lowercase         = torch.Size((1, 1_0_0_0)				)
      self.assertEqual(outputs.logits.shape							,			A__				)

      lowercase         = torch.tensor([1.0_4_4_8, -0.3_7_4_5, -1.8_3_1_7]				).to(A__				)

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











'''simple docstring'''

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


__lowerCamelCase : Optional[Any]											=			logging.get_logger(__name__)

__lowerCamelCase : Optional[Any]											=			{
    "facebook/vit-mae-base": "https://huggingface.co/facebook/vit-mae-base/resolve/main/config.json",
    # See all ViT MAE models at https://huggingface.co/models?filter=vit-mae
}

class       UpperCAmelCase (						_lowercase  ):
		UpperCAmelCase					:	Union[str, Any] 			=		'''vit_mae'''




		def __init__(self      :	Dict							,			A__      :	Tuple=7_6_8							,			A__      :	str=1_2							,			A__      :	Dict=1_2							,			A__      :	Optional[Any]=3_0_7_2							,			A__      :	Union[str, Any]="gelu"							,			A__      :	Optional[int]=0.0							,			A__      :	str=0.0							,			A__      :	Optional[Any]=0.0_2							,			A__      :	Optional[Any]=1e-12							,			A__      :	Tuple=2_2_4							,			A__      :	str=1_6							,			A__      :	Optional[int]=3							,			A__      :	List[Any]=True							,			A__      :	Dict=1_6							,			A__      :	Dict=5_1_2							,			A__      :	Dict=8							,			A__      :	Union[str, Any]=2_0_4_8							,			A__      :	str=0.7_5							,			A__      :	str=False							,			**A__      :	Any							,			)					-> Union[str, Any]:
						super().__init__(**A__				)

						lowercase         = hidden_size
						lowercase         = num_hidden_layers
						lowercase         = num_attention_heads
						lowercase         = intermediate_size
						lowercase         = hidden_act
						lowercase         = hidden_dropout_prob
						lowercase         = attention_probs_dropout_prob
						lowercase         = initializer_range
						lowercase         = layer_norm_eps
						lowercase         = image_size
						lowercase         = patch_size
						lowercase         = num_channels
						lowercase         = qkv_bias
						lowercase         = decoder_num_attention_heads
						lowercase         = decoder_hidden_size
						lowercase         = decoder_num_hidden_layers
						lowercase         = decoder_intermediate_size
						lowercase         = mask_ratio
						lowercase         = norm_pix_loss




