
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     __A:




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


								'''simple docstring'''




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

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




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


								'''simple docstring'''




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

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

								__a    							=					self.get_config()

								return config, pixel_values, labels, pixel_labels




				def        SCREAMING_SNAKE_CASE_	(  self       )					->   List[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=_snake_case	,  initializer_range=self.initializer_range	,  out_indices=self.out_indices	,  )




				def        SCREAMING_SNAKE_CASE_	(  self	,  _snake_case	,  _snake_case	,  _snake_case	,  _snake_case       )					->   List[Any]:


								'''simple docstring'''




								__a    							=					DataaVecVisionModel(config=_snake_case       )
								model.to(_snake_case       )
								model.eval()
								__a    							=					model(_snake_case       )
								# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
								__a    							=					(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        SCREAMING_SNAKE_CASE_	(  self	,  _snake_case	,  _snake_case	,  _snake_case	,  _snake_case       )					->   Tuple:


								'''simple docstring'''




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




				def        SCREAMING_SNAKE_CASE_	(  self	,  _snake_case	,  _snake_case	,  _snake_case	,  _snake_case       )					->   int:


								'''simple docstring'''




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




				def        SCREAMING_SNAKE_CASE_	(  self       )					->   Tuple:


								'''simple docstring'''




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







@require_torch
class     __A(				a						,     a						,     unittest.TestCase    ):
				snake_case_		     =					(
				    (DataaVecVisionModel, DataaVecVisionForImageClassification, DataaVecVisionForSemanticSegmentation)
				    if is_torch_available()
				    else ()
				)
				snake_case_		     =					(
				    {
				        '''feature-extraction''': DataaVecVisionModel,
				        '''image-classification''': DataaVecVisionForImageClassification,
				        '''image-segmentation''': DataaVecVisionForSemanticSegmentation,
				    }
				    if is_torch_available()
				    else {}
				)

				snake_case_		     =					False
				snake_case_		     =					False
				snake_case_		     =					False




				def        SCREAMING_SNAKE_CASE_	(  self       )					->   Dict:


								'''simple docstring'''




								__a    							=					DataaVecVisionModelTester(self       )
								__a    							=					ConfigTester(
								    self	,  config_class=_snake_case	,  has_text_modality=_snake_case	,  hidden_size=37       )




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


								'''simple docstring'''




								self.config_tester.run_common_tests()




				def        SCREAMING_SNAKE_CASE_	(  self       )					->   List[Any]:


								'''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        SCREAMING_SNAKE_CASE_	(  self       )					->   Dict:


								'''simple docstring'''




								pass




				def        SCREAMING_SNAKE_CASE_	(  self       )					->   Any:


								'''simple docstring'''




								__a							,  __a    							=					self.model_tester.prepare_config_and_inputs_for_common()

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




				def        SCREAMING_SNAKE_CASE_	(  self       )					->   List[Any]:


								'''simple docstring'''




								__a							,  __a    							=					self.model_tester.prepare_config_and_inputs_for_common()

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

												__a    							=					['''pixel_values''']
												self.assertListEqual(arg_names[:1]	,  _snake_case       )




				def        SCREAMING_SNAKE_CASE_	(  self       )					->   List[str]:


								'''simple docstring'''




								__a    							=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_snake_case       )




				def        SCREAMING_SNAKE_CASE_	(  self       )					->   List[Any]:


								'''simple docstring'''




								__a    							=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_segmentation(*_snake_case       )




				def        SCREAMING_SNAKE_CASE_	(  self       )					->   str:


								'''simple docstring'''




								if not self.model_tester.is_training:
												return

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

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

												__a    							=					model_class(_snake_case       )
												model.to(_snake_case       )
												model.train()
												__a    							=					self._prepare_for_class(_snake_case	,  _snake_case	,  return_labels=_snake_case       )
												__a    							=					model(**_snake_case       ).loss
												loss.backward()




				def        SCREAMING_SNAKE_CASE_	(  self       )					->   Any:


								'''simple docstring'''




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

								__a    							=					False
								__a    							=					True

								for model_class in self.all_model_classes:
												if model_class in [*get_values(_snake_case       )] 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":
																__a							,  __a							,  __a							,  __a    							=					inputs_dict['''pixel_values'''].shape
																__a    							=					torch.zeros(
																    [self.model_tester.batch_size, height, width]	,  device=_snake_case       ).long()
												__a    							=					model_class(_snake_case       )
												model.gradient_checkpointing_enable()
												model.to(_snake_case       )
												model.train()
												__a    							=					self._prepare_for_class(_snake_case	,  _snake_case	,  return_labels=_snake_case       )
												__a    							=					model(**_snake_case       ).loss
												loss.backward()




				def        SCREAMING_SNAKE_CASE_	(  self       )					->   List[Any]:


								'''simple docstring'''




								__a							,  __a    							=					self.model_tester.prepare_config_and_inputs_for_common()

								__a    							=					_config_zero_init(_snake_case       )
								for model_class in self.all_model_classes:
												__a    							=					model_class(config=_snake_case       )
												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        SCREAMING_SNAKE_CASE_	(  self	,  _snake_case	,  _snake_case	,  _snake_case	,  _snake_case=2E-4	,  _snake_case="outputs"	,  _snake_case=None       )					->   Tuple:


								'''simple docstring'''




								super().check_pt_tf_outputs(_snake_case	,  _snake_case	,  _snake_case	,  _snake_case	,  _snake_case	,  _snake_case       )




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


								'''simple docstring'''




								__a    							=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*_snake_case       )




				@slow
				def        SCREAMING_SNAKE_CASE_	(  self       )					->   List[str]:


								'''simple docstring'''




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



def 					__lowerCAmelCase    (    ) ->						str:
				__a    							=					Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''						)
				return image







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




				@cached_property
				def        SCREAMING_SNAKE_CASE_	(  self       )					->   List[Any]:


								'''simple docstring'''




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




				@slow
				def        SCREAMING_SNAKE_CASE_	(  self       )					->   int:


								'''simple docstring'''




								__a    							=					DataaVecVisionForImageClassification.from_pretrained('''facebook/data2vec-vision-base-ft1k'''       ).to(
								    _snake_case       )

								__a    							=					self.default_image_processor
								__a    							=					prepare_img()
								__a    							=					image_processor(images=_snake_case	,  return_tensors='''pt'''       ).to(_snake_case       )

								# forward pass
								with torch.no_grad():
												__a    							=					model(**_snake_case       )
								__a    							=					outputs.logits

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

								__a    							=					torch.tensor([0.3277, -0.1395, 0.0911]       ).to(_snake_case       )

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

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








from itertools import compress, repeat
from math import ceil, sqrt


def 					__lowerCAmelCase    (    a__						) ->						list[int]:

				if num <= 2:
								return []
				if num == 3:
								return [2]

				# Odd sieve for numbers in range [3, num - 1]
				__a    							=					bytearray(B'''\x01'''						) * ((num >> 1) - 1)

				for i in range(3		,					int(sqrt(a__						)						) + 1		,					2						):
								if sieve[(i >> 1) - 1]:
												__a    							=					i**2
												__a    							=					repeat(
												    0		,					ceil((num - i_squared) / (i << 1)						)						)

				return [2] + list(compress(range(3		,					a__		,					2						)		,					a__						)						)


if __name__ == "__main__":
					import doctest

					doctest.testmod()