


"""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       lowerCamelCase_		:





				"""simple docstring"""







				def __init__(    self							,		UpperCAmelCase__							,		UpperCAmelCase__=100							,		UpperCAmelCase__=13							,		UpperCAmelCase__=30							,		UpperCAmelCase__=2							,		UpperCAmelCase__=3							,		UpperCAmelCase__=True							,		UpperCAmelCase__=True							,		UpperCAmelCase__=32							,		UpperCAmelCase__=4							,		UpperCAmelCase__=4							,		UpperCAmelCase__=37							,		UpperCAmelCase__="gelu"							,		UpperCAmelCase__=0.1							,		UpperCAmelCase__=0.1							,		UpperCAmelCase__=10							,		UpperCAmelCase__=0.02							,		UpperCAmelCase__=3							,		UpperCAmelCase__=None							,		UpperCAmelCase__=[0, 1, 2, 3]							,		):
								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 	lowerCAmelCase__		(    self  ):
								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 	lowerCAmelCase__		(    self  ):
								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=_lowercase							,		initializer_range=self.initializer_range							,		out_indices=self.out_indices							,		)



				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__  ):
								SCREAMING_SNAKE_CASE__  				=					DataaVecVisionModel(config=_lowercase  )
								model.to(_lowercase  )
								model.eval()
								SCREAMING_SNAKE_CASE__  				=					model(_lowercase  )
								# 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 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__  ):
								SCREAMING_SNAKE_CASE__  				=					self.type_sequence_label_size
								SCREAMING_SNAKE_CASE__  				=					DataaVecVisionForImageClassification(_lowercase  )
								model.to(_lowercase  )
								model.eval()
								SCREAMING_SNAKE_CASE__  				=					model(_lowercase							,		labels=_lowercase  )
								self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.type_sequence_label_size)  )



				def 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__  ):
								SCREAMING_SNAKE_CASE__  				=					self.num_labels
								SCREAMING_SNAKE_CASE__  				=					DataaVecVisionForSemanticSegmentation(_lowercase  )
								model.to(_lowercase  )
								model.eval()
								SCREAMING_SNAKE_CASE__  				=					model(_lowercase  )
								self.parent.assertEqual(
								    result.logits.shape							,		(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)  )
								SCREAMING_SNAKE_CASE__  				=					model(_lowercase							,		labels=_lowercase  )
								self.parent.assertEqual(
								    result.logits.shape							,		(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)  )






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


@require_torch
class       lowerCamelCase_		(      UpperCAmelCase_							,	UpperCAmelCase_							,	unittest.TestCase ):





				"""simple docstring"""







				_lowerCAmelCase   :   Dict   				=       (
				    (DataaVecVisionModel, DataaVecVisionForImageClassification, DataaVecVisionForSemanticSegmentation)
				    if is_torch_available()
				    else ()
				)
				_lowerCAmelCase   :   Optional[Any]   				=       (
				    {
				        'feature-extraction': DataaVecVisionModel,
				        'image-classification': DataaVecVisionForImageClassification,
				        'image-segmentation': DataaVecVisionForSemanticSegmentation,
				    }
				    if is_torch_available()
				    else {}
				)

				_lowerCAmelCase   :   int   				=       False
				_lowerCAmelCase   :   Optional[Any]   				=       False
				_lowerCAmelCase   :   Tuple   				=       False



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					DataaVecVisionModelTester(self  )
								SCREAMING_SNAKE_CASE__  				=					ConfigTester(
								    self							,		config_class=_lowercase							,		has_text_modality=_lowercase							,		hidden_size=37  )



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



				def 	lowerCAmelCase__		(    self  ):
								# Data2VecVision does not use inputs_embeds
								pass



				@require_torch_multi_gpu
				@unittest.skip(
				    reason="Data2VecVision has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`"  )
				def 	lowerCAmelCase__		(    self  ):
								pass



				def 	lowerCAmelCase__		(    self  ):
								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(_lowercase  )
												self.assertIsInstance(model.get_input_embeddings()							,		(nn.Module)  )
												SCREAMING_SNAKE_CASE__  				=					model.get_output_embeddings()
												self.assertTrue(x is None or isinstance(_lowercase							,		nn.Linear  )  )



				def 	lowerCAmelCase__		(    self  ):
								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(_lowercase  )
												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]							,		_lowercase  )



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_lowercase  )



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_segmentation(*_lowercase  )



				def 	lowerCAmelCase__		(    self  ):
								if not self.model_tester.is_training:
												return

								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(_lowercase  )]:
																continue

												SCREAMING_SNAKE_CASE__  				=					model_class(_lowercase  )
												model.to(_lowercase  )
												model.train()
												SCREAMING_SNAKE_CASE__  				=					self._prepare_for_class(_lowercase							,		_lowercase							,		return_labels=_lowercase  )
												SCREAMING_SNAKE_CASE__  				=					model(**_lowercase  ).loss
												loss.backward()



				def 	lowerCAmelCase__		(    self  ):
								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(_lowercase  )] 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__  				=					inputs_dict['pixel_values'].shape
																SCREAMING_SNAKE_CASE__  				=					torch.zeros(
																    [self.model_tester.batch_size, height, width]							,		device=_lowercase  ).long()
												SCREAMING_SNAKE_CASE__  				=					model_class(_lowercase  )
												model.gradient_checkpointing_enable()
												model.to(_lowercase  )
												model.train()
												SCREAMING_SNAKE_CASE__  				=					self._prepare_for_class(_lowercase							,		_lowercase							,		return_labels=_lowercase  )
												SCREAMING_SNAKE_CASE__  				=					model(**_lowercase  ).loss
												loss.backward()



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					self.model_tester.prepare_config_and_inputs_for_common()

								SCREAMING_SNAKE_CASE__  				=					_config_zero_init(_lowercase  )
								for model_class in self.all_model_classes:
												SCREAMING_SNAKE_CASE__  				=					model_class(config=_lowercase  )
												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 	lowerCAmelCase__		(    self							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__							,		UpperCAmelCase__=2e-4							,		UpperCAmelCase__="outputs"							,		UpperCAmelCase__=None  ):
								# We override with a slightly higher tol value, as semseg models tend to diverge a bit more
								super().check_pt_tf_outputs(_lowercase							,		_lowercase							,		_lowercase							,		_lowercase							,		_lowercase							,		_lowercase  )



				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*_lowercase  )






				@slow
				def 	lowerCAmelCase__		(    self  ):
								for model_name in DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												SCREAMING_SNAKE_CASE__  				=					DataaVecVisionModel.from_pretrained(_lowercase  )
												self.assertIsNotNone(_lowercase  )


def __lowercase  (   ):
				SCREAMING_SNAKE_CASE__  				=					Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"	)
				return image


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





				"""simple docstring"""







				@cached_property
				def 	lowerCAmelCase__		(    self  ):
								return (
								    BeitImageProcessor.from_pretrained("facebook/data2vec-vision-base-ft1k"  ) if is_vision_available() else None
								)






				@slow
				def 	lowerCAmelCase__		(    self  ):
								SCREAMING_SNAKE_CASE__  				=					DataaVecVisionForImageClassification.from_pretrained("facebook/data2vec-vision-base-ft1k"  ).to(
								    _lowercase  )

								SCREAMING_SNAKE_CASE__  				=					self.default_image_processor
								SCREAMING_SNAKE_CASE__  				=					prepare_img()
								SCREAMING_SNAKE_CASE__  				=					image_processor(images=_lowercase							,		return_tensors="pt"  ).to(_lowercase  )

								# forward pass
								with torch.no_grad():
												SCREAMING_SNAKE_CASE__  				=					model(**_lowercase  )
								SCREAMING_SNAKE_CASE__  				=					outputs.logits

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

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

								self.assertTrue(torch.allclose(logits[0, :3]							,		_lowercase							,		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()							,		_lowercase  )








'''simple docstring'''

import inspect
import warnings

import torch

from .state import AcceleratorState, GradientState
from .utils import DistributedType, honor_type, is_tpu_available


if is_tpu_available(check_device=False):
						import torch_xla.core.xla_model as xm






def 						_UpperCAmelCase     (	a		:						Tuple    ,      a		:						int				)      ->      List[str]:





						"""simple docstring"""
						if isinstance(a    ,      (list, tuple)				):
												return honor_type(a    ,      (move_to_device(a    ,      a				) for t in state)				)
						elif isinstance(a    ,      a				):
												return type(a				)({k: move_to_device(a    ,      a				) for k, v in state.items()}				)
						elif isinstance(a    ,      torch.Tensor				):
												return state.to(a				)
						return state






class 				__magic_name__  (						torch.optim.Optimizer					):







	"""simple docstring"""



	def __init__(				self	,   _lowercase	,   _lowercase=True	,   _lowercase=None							)    ->	int:
							lowercase_   : List[str]   	=      optimizer
							lowercase_   : Optional[int]   	=      scaler
							lowercase_   : Union[str, Any]   	=      AcceleratorState()
							lowercase_   : Optional[Any]   	=      GradientState()
							lowercase_   : Union[str, Any]   	=      device_placement
							lowercase_   : List[str]   	=      False
							lowercase_   : Optional[Any]   	=      None

							# Handle device placement
							if device_placement:
													lowercase_   : str   	=      self.optimizer.state_dict()
													if self.accelerator_state.distributed_type == DistributedType.TPU:
																			xm.send_cpu_data_to_device(_lowercase	,   self.accelerator_state.device							)
													else:
																			lowercase_   : List[str]   	=      move_to_device(_lowercase	,   self.accelerator_state.device							)
													self.optimizer.load_state_dict(_lowercase							)
	@property
	def 							lowerCamelCase__		(				self							)    ->	Any:
							return self.optimizer.state
	@state.setter
	def 							lowerCamelCase__		(				self	,   _lowercase							)    ->	Union[str, Any]:
							lowercase_   : Optional[Any]   	=      state
	@property
	def 							lowerCamelCase__		(				self							)    ->	Any:
							return self.optimizer.param_groups
	@param_groups.setter
	def 							lowerCamelCase__		(				self	,   _lowercase							)    ->	Optional[Any]:
							lowercase_   : Dict   	=      param_groups
	@property
	def 							lowerCamelCase__		(				self							)    ->	int:
							return self.optimizer.defaults
	@defaults.setter
	def 							lowerCamelCase__		(				self	,   _lowercase							)    ->	Any:
							lowercase_   : List[Any]   	=      defaults
	def 							lowerCamelCase__		(				self	,   _lowercase							)    ->	Tuple:
							self.optimizer.add_param_group(_lowercase							)
	def 							lowerCamelCase__		(				self	,   _lowercase							)    ->	Any:
							if self.accelerator_state.distributed_type == DistributedType.TPU and self.device_placement:
													xm.send_cpu_data_to_device(_lowercase	,   self.accelerator_state.device							)
							self.optimizer.load_state_dict(_lowercase							)
	def 							lowerCamelCase__		(				self							)    ->	Any:
							return self.optimizer.state_dict()
	def 							lowerCamelCase__		(				self	,   _lowercase=None							)    ->	Optional[Any]:
							if self.gradient_state.sync_gradients:
													lowercase_   : Tuple   	=      'set_to_none' in inspect.signature(self.optimizer.zero_grad							).parameters
													if accept_arg:
																			if set_to_none is None:
																									lowercase_   : Optional[int]   	=      False
																			self.optimizer.zero_grad(set_to_none=_lowercase							)
													else:
																			if set_to_none is not None:
																									raise ValueError('`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer.'							)
																			self.optimizer.zero_grad()
	def 							lowerCamelCase__		(				self	,   _lowercase=None							)    ->	Optional[Any]:
							if self.gradient_state.sync_gradients:
													if self.accelerator_state.distributed_type == DistributedType.TPU:
																			lowercase_   : Union[str, Any]   	=      {'closure': closure} if closure is not None else {}
																			xm.optimizer_step(self.optimizer	,   optimizer_args=_lowercase							)
													elif self.scaler is not None:
																			lowercase_   : Optional[Any]   	=      False
																			if self._last_scale is None:
																									# `get_scale` is an async operation requiring full synchronization
																									# on CPU and GPUs before finishing. As a result, we store away
																									# the prior one to reduce the call overhead
																									lowercase_   : Tuple   	=      self.scaler.get_scale()
																									lowercase_   : int   	=      True
																			self.scaler.step(self.optimizer	,   _lowercase							)
																			self.scaler.update()
																			lowercase_   : List[str]   	=      self.scaler.get_scale()
																			if not new_scale:
																									# If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
																									lowercase_   : Optional[int]   	=      scale_after < self._last_scale
																			lowercase_   : List[str]   	=      scale_after

													else:
																			self.optimizer.step(_lowercase							)
	def 							lowerCamelCase__		(				self	,   _lowercase							)    ->	Optional[int]:
							for param_group in self.optimizer.param_groups:
													lowercase_   : Dict   	=      [parameters_map.get(_lowercase	,   _lowercase							) for p in param_group['params']]
	@property
	def 							lowerCamelCase__		(				self							)    ->	List[Any]:
							warnings.warn(
							    'The `is_overflow` property is deprecated and will be removed in version 1.0 of Accelerate use '
							    '`optimizer.step_was_skipped` instead.'	,   _lowercase	,   )
							return self._is_overflow
	@property
	def 							lowerCamelCase__		(				self							)    ->	str:
							return self._is_overflow
	def __getstate__(				self							)    ->	str:
							return self.__dict__.copy()







	def __setstate__(				self	,   _lowercase							)    ->	Union[str, Any]:
							self.__dict__.update(_lowercase							)
