
"""simple docstring"""




import inspect
import unittest

from transformers import DataaVecVisionConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_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, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch
	from torch import nn

	from transformers import (
	    MODEL_MAPPING,
	    DataaVecVisionForImageClassification,
	    DataaVecVisionForSemanticSegmentation,
	    DataaVecVisionModel,
	)
	from transformers.models.dataavec.modeling_dataavec_vision import DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from PIL import Image

	from transformers import BeitImageProcessor




class 	UpperCamelCase      :

				"""simple docstring"""



				def __init__(							self						,      lowercase__						,      lowercase__=100						,      lowercase__=13						,      lowercase__=30						,      lowercase__=2						,      lowercase__=3						,      lowercase__=True						,      lowercase__=True						,      lowercase__=32						,      lowercase__=4						,      lowercase__=4						,      lowercase__=37						,      lowercase__="gelu"						,      lowercase__=0.1						,      lowercase__=0.1						,      lowercase__=10						,      lowercase__=0.02						,      lowercase__=3						,      lowercase__=None						,      lowercase__=[0, 1, 2, 3]						,      )       ->       Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =      parent
									SCREAMING_SNAKE_CASE           =      100
									SCREAMING_SNAKE_CASE           =      batch_size
									SCREAMING_SNAKE_CASE           =      image_size
									SCREAMING_SNAKE_CASE           =      patch_size
									SCREAMING_SNAKE_CASE           =      num_channels
									SCREAMING_SNAKE_CASE           =      is_training
									SCREAMING_SNAKE_CASE           =      use_labels
									SCREAMING_SNAKE_CASE           =      hidden_size
									SCREAMING_SNAKE_CASE           =      num_hidden_layers
									SCREAMING_SNAKE_CASE           =      num_attention_heads
									SCREAMING_SNAKE_CASE           =      intermediate_size
									SCREAMING_SNAKE_CASE           =      hidden_act
									SCREAMING_SNAKE_CASE           =      hidden_dropout_prob
									SCREAMING_SNAKE_CASE           =      attention_probs_dropout_prob
									SCREAMING_SNAKE_CASE           =      type_sequence_label_size
									SCREAMING_SNAKE_CASE           =      initializer_range
									SCREAMING_SNAKE_CASE           =      scope
									SCREAMING_SNAKE_CASE           =      out_indices
									SCREAMING_SNAKE_CASE           =      num_labels

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


				def 				A   (							self     )       ->       Dict:
									"""simple docstring"""







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

									SCREAMING_SNAKE_CASE           =      None
									SCREAMING_SNAKE_CASE           =      None
									if self.use_labels:
														SCREAMING_SNAKE_CASE           =      ids_tensor([self.batch_size]						,      self.type_sequence_label_size     )
														SCREAMING_SNAKE_CASE           =      ids_tensor([self.batch_size, self.image_size, self.image_size]						,      self.num_labels     )

									SCREAMING_SNAKE_CASE           =      self.get_config()

									return config, pixel_values, labels, pixel_labels


				def 				A   (							self     )       ->       Union[str, Any]:
									"""simple docstring"""







									return DataaVecVisionConfig(
									    vocab_size=self.vocab_size						,      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						,      out_indices=self.out_indices						,      )


				def 				A   (							self						,      lowercase__						,      lowercase__						,      lowercase__						,      lowercase__     )       ->       Any:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =      DataaVecVisionModel(config=_lowerCAmelCase     )
									model.to(_lowerCAmelCase     )
									model.eval()
									SCREAMING_SNAKE_CASE           =      model(_lowerCAmelCase     )
									# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
									SCREAMING_SNAKE_CASE           =      (self.image_size // self.patch_size) ** 2
									self.parent.assertEqual(result.last_hidden_state.shape						,      (self.batch_size, num_patches + 1, self.hidden_size)     )


				def 				A   (							self						,      lowercase__						,      lowercase__						,      lowercase__						,      lowercase__     )       ->       Dict:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =      self.type_sequence_label_size
									SCREAMING_SNAKE_CASE           =      DataaVecVisionForImageClassification(_lowerCAmelCase     )
									model.to(_lowerCAmelCase     )
									model.eval()
									SCREAMING_SNAKE_CASE           =      model(_lowerCAmelCase						,      labels=_lowerCAmelCase     )
									self.parent.assertEqual(result.logits.shape						,      (self.batch_size, self.type_sequence_label_size)     )


				def 				A   (							self						,      lowercase__						,      lowercase__						,      lowercase__						,      lowercase__     )       ->       Tuple:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =      self.num_labels
									SCREAMING_SNAKE_CASE           =      DataaVecVisionForSemanticSegmentation(_lowerCAmelCase     )
									model.to(_lowerCAmelCase     )
									model.eval()
									SCREAMING_SNAKE_CASE           =      model(_lowerCAmelCase     )
									self.parent.assertEqual(
									    result.logits.shape						,      (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)     )
									SCREAMING_SNAKE_CASE           =      model(_lowerCAmelCase						,      labels=_lowerCAmelCase     )
									self.parent.assertEqual(
									    result.logits.shape						,      (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)     )




				def 				A   (							self     )       ->       Tuple:
									"""simple docstring"""







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






@require_torch
class 	UpperCamelCase      (			_UpperCAmelCase    , _UpperCAmelCase    , unittest.TestCase ):

				"""simple docstring"""



				UpperCAmelCase_							: Union[str, Any]		    =							(
				    (DataaVecVisionModel, DataaVecVisionForImageClassification, DataaVecVisionForSemanticSegmentation)
				    if is_torch_available()
				    else ()
				)
				UpperCAmelCase_							: Optional[int]		    =							(
				    {
				        "feature-extraction": DataaVecVisionModel,
				        "image-classification": DataaVecVisionForImageClassification,
				        "image-segmentation": DataaVecVisionForSemanticSegmentation,
				    }
				    if is_torch_available()
				    else {}
				)

				UpperCAmelCase_							: Optional[int]		    =							False
				UpperCAmelCase_							: Optional[int]		    =							False
				UpperCAmelCase_							: List[str]		    =							False


				def 				A   (							self     )       ->       Union[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =      DataaVecVisionModelTester(self     )
									SCREAMING_SNAKE_CASE           =      ConfigTester(
									    self						,      config_class=_lowerCAmelCase						,      has_text_modality=_lowerCAmelCase						,      hidden_size=37     )


				def 				A   (							self     )       ->       Any:
									"""simple docstring"""







									self.config_tester.run_common_tests()


				def 				A   (							self     )       ->       int:
									"""simple docstring"""







									pass


				@require_torch_multi_gpu
				@unittest.skip(
				    reason='Data2VecVision has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'     )
				def 				A   (							self     )       ->       Union[str, Any]:
									"""simple docstring"""







									pass


				def 				A   (							self     )       ->       str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE           =      self.model_tester.prepare_config_and_inputs_for_common()

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


				def 				A   (							self     )       ->       Any:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE           =      self.model_tester.prepare_config_and_inputs_for_common()

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

														SCREAMING_SNAKE_CASE           =      ['pixel_values']
														self.assertListEqual(arg_names[:1]						,      _lowerCAmelCase     )


				def 				A   (							self     )       ->       Dict:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_lowerCAmelCase     )


				def 				A   (							self     )       ->       Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_segmentation(*_lowerCAmelCase     )


				def 				A   (							self     )       ->       Union[str, Any]:
									"""simple docstring"""







									if not self.model_tester.is_training:
														return

									SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE           =      self.model_tester.prepare_config_and_inputs_for_common()
									SCREAMING_SNAKE_CASE           =      True

									for model_class in self.all_model_classes:
														if model_class in [*get_values(_lowerCAmelCase     )]:
																			continue

														SCREAMING_SNAKE_CASE           =      model_class(_lowerCAmelCase     )
														model.to(_lowerCAmelCase     )
														model.train()
														SCREAMING_SNAKE_CASE           =      self._prepare_for_class(_lowerCAmelCase						,      _lowerCAmelCase						,      return_labels=_lowerCAmelCase     )
														SCREAMING_SNAKE_CASE           =      model(**_lowerCAmelCase     ).loss
														loss.backward()


				def 				A   (							self     )       ->       Any:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE           =      self.model_tester.prepare_config_and_inputs_for_common()
									if not self.model_tester.is_training:
														return

									SCREAMING_SNAKE_CASE           =      False
									SCREAMING_SNAKE_CASE           =      True

									for model_class in self.all_model_classes:
														if model_class in [*get_values(_lowerCAmelCase     )] or not model_class.supports_gradient_checkpointing:
																			continue
														# TODO: remove the following 3 lines once we have a MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING
														# this can then be incorporated into _prepare_for_class in test_modeling_common.py
														elif model_class.__name__ == "Data2VecVisionForSemanticSegmentation":
																			SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE           =      inputs_dict['pixel_values'].shape
																			SCREAMING_SNAKE_CASE           =      torch.zeros(
																			    [self.model_tester.batch_size, height, width]						,      device=_lowerCAmelCase     ).long()
														SCREAMING_SNAKE_CASE           =      model_class(_lowerCAmelCase     )
														model.gradient_checkpointing_enable()
														model.to(_lowerCAmelCase     )
														model.train()
														SCREAMING_SNAKE_CASE           =      self._prepare_for_class(_lowerCAmelCase						,      _lowerCAmelCase						,      return_labels=_lowerCAmelCase     )
														SCREAMING_SNAKE_CASE           =      model(**_lowerCAmelCase     ).loss
														loss.backward()


				def 				A   (							self     )       ->       int:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE           =      self.model_tester.prepare_config_and_inputs_for_common()

									SCREAMING_SNAKE_CASE           =      _config_zero_init(_lowerCAmelCase     )
									for model_class in self.all_model_classes:
														SCREAMING_SNAKE_CASE           =      model_class(config=_lowerCAmelCase     )
														for name, param in model.named_parameters():
																			# we skip lambda parameters as these require special initial values
																			# determined by config.layer_scale_init_value
																			if "lambda" in name:
																								continue
																			if param.requires_grad:
																								self.assertIn(
																								    ((param.data.mean() * 1E9).round() / 1E9).item()						,      [0.0, 1.0]						,      msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''						,      )


				def 				A   (							self						,      lowercase__						,      lowercase__						,      lowercase__						,      lowercase__=2E-4						,      lowercase__="outputs"						,      lowercase__=None     )       ->       Union[str, Any]:
									"""simple docstring"""







									super().check_pt_tf_outputs(_lowerCAmelCase						,      _lowerCAmelCase						,      _lowerCAmelCase						,      _lowerCAmelCase						,      _lowerCAmelCase						,      _lowerCAmelCase     )


				def 				A   (							self     )       ->       Tuple:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*_lowerCAmelCase     )




				@slow
				def 				A   (							self     )       ->       List[Any]:
									"""simple docstring"""







									for model_name in DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														SCREAMING_SNAKE_CASE           =      DataaVecVisionModel.from_pretrained(_lowerCAmelCase     )
														self.assertIsNotNone(_lowerCAmelCase     )






def     UpperCamelCase_    (						):
					SCREAMING_SNAKE_CASE           =      Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'     )
					return image




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

				"""simple docstring"""



				@cached_property
				def 				A   (							self     )       ->       Tuple:
									"""simple docstring"""







									return (
									    BeitImageProcessor.from_pretrained('facebook/data2vec-vision-base-ft1k'     ) if is_vision_available() else None
									)




				@slow
				def 				A   (							self     )       ->       str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =      DataaVecVisionForImageClassification.from_pretrained('facebook/data2vec-vision-base-ft1k'     ).to(
									    _lowerCAmelCase     )

									SCREAMING_SNAKE_CASE           =      self.default_image_processor
									SCREAMING_SNAKE_CASE           =      prepare_img()
									SCREAMING_SNAKE_CASE           =      image_processor(images=_lowerCAmelCase						,      return_tensors='pt'     ).to(_lowerCAmelCase     )

									# forward pass
									with torch.no_grad():
														SCREAMING_SNAKE_CASE           =      model(**_lowerCAmelCase     )
									SCREAMING_SNAKE_CASE           =      outputs.logits

									# verify the logits
									SCREAMING_SNAKE_CASE           =      torch.Size((1, 1000)     )
									self.assertEqual(logits.shape						,      _lowerCAmelCase     )

									SCREAMING_SNAKE_CASE           =      torch.tensor([0.3_277, -0.1_395, 0.0_911]     ).to(_lowerCAmelCase     )

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

									SCREAMING_SNAKE_CASE           =      [model.config.labelaid[i] for i in ['remote control, remote', 'tabby, tabby cat']]
									self.assertEqual(logits[0].topk(2     ).indices.cpu().tolist()						,      _lowerCAmelCase     )




from ..utils import DummyObject, requires_backends

class      __snake_case				(   metaclass=_UpperCAmelCase):
					SCREAMING_SNAKE_CASE_		:       str								=			["torch", "transformers", "k_diffusion"]
					def __init__(						self       :						Tuple	,      *_lowerCAmelCase       :						Union[str, Any]	,      **_lowerCAmelCase       :						Union[str, Any]  )							->					Dict:



											'''simple docstring'''





											requires_backends(self	,      ['torch', 'transformers', 'k_diffusion']  )
					@classmethod
					def 	a			(						cls       :						Optional[int]	,      *_lowerCAmelCase       :						Any	,      **_lowerCAmelCase       :						Optional[Any]  )							->					Optional[Any]:



											'''simple docstring'''





											requires_backends(cls	,      ['torch', 'transformers', 'k_diffusion']  )







					@classmethod
					def 	a			(						cls       :						Union[str, Any]	,      *_lowerCAmelCase       :						Tuple	,      **_lowerCAmelCase       :						str  )							->					Tuple:



											'''simple docstring'''





											requires_backends(cls	,      ['torch', 'transformers', 'k_diffusion']  )

