
'''simple docstring'''
from numpy import exp, pi, sqrt







def 		A_(     A						:		Any		,     A						:		float = 0.0		,     A						:		float = 1.0):
							return 1 / sqrt(2 * pi * sigma**2) * exp(-((x - mu) ** 2) / (2 * sigma**2))


if __name__ == "__main__":
						import doctest

						doctest.testmod()






'''simple docstring'''
import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
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,
						    BeitForImageClassification,
						    BeitForMaskedImageModeling,
						    BeitForSemanticSegmentation,
						    BeitModel,
						)
						from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						import PIL
						from PIL import Image

						from transformers import BeitImageProcessor


class 			SCREAMING_SNAKE_CASE__      :







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







											'''simple docstring'''






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

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







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







											'''simple docstring'''






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

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

											UpperCamelCase   			=     self.get_config()

											return config, pixel_values, labels, pixel_labels







				def 						UpperCAmelCase_    (       self     )->		Dict:







											'''simple docstring'''






											return BeitConfig(
											    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=A_						,				initializer_range=self.initializer_range						,				out_indices=self.out_indices						,				)







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				A_						,				A_     )->		List[str]:







											'''simple docstring'''






											UpperCamelCase   			=     BeitModel(config=A_     )
											model.to(A_     )
											model.eval()
											UpperCamelCase   			=     model(A_     )
											self.parent.assertEqual(result.last_hidden_state.shape						,				(self.batch_size, self.seq_length, self.hidden_size)     )







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				A_						,				A_     )->		Any:







											'''simple docstring'''






											UpperCamelCase   			=     BeitForMaskedImageModeling(config=A_     )
											model.to(A_     )
											model.eval()
											UpperCamelCase   			=     model(A_     )
											self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.seq_length - 1, self.vocab_size)     )







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				A_						,				A_     )->		Optional[int]:







											'''simple docstring'''






											UpperCamelCase   			=     self.type_sequence_label_size
											UpperCamelCase   			=     BeitForImageClassification(A_     )
											model.to(A_     )
											model.eval()
											UpperCamelCase   			=     model(A_						,				labels=A_     )
											self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.type_sequence_label_size)     )

											# test greyscale images
											UpperCamelCase   			=     1
											UpperCamelCase   			=     BeitForImageClassification(A_     )
											model.to(A_     )
											model.eval()

											UpperCamelCase   			=     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
											UpperCamelCase   			=     model(A_						,				labels=A_     )
											self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.type_sequence_label_size)     )







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				A_						,				A_     )->		Optional[Any]:







											'''simple docstring'''






											UpperCamelCase   			=     self.num_labels
											UpperCamelCase   			=     BeitForSemanticSegmentation(A_     )
											model.to(A_     )
											model.eval()
											UpperCamelCase   			=     model(A_     )
											self.parent.assertEqual(
											    result.logits.shape						,				(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)     )
											UpperCamelCase   			=     model(A_						,				labels=A_     )
											self.parent.assertEqual(
											    result.logits.shape						,				(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)     )







				def 						UpperCAmelCase_    (       self     )->		int:







											'''simple docstring'''






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


@require_torch
class 			SCREAMING_SNAKE_CASE__      (     snake_case_							,       snake_case_							,       unittest.TestCase):
				lowerCAmelCase_     			=							(
				    (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
				    if is_torch_available()
				    else ()
				)
				lowerCAmelCase_     			=							(
				    {
				        """feature-extraction""": BeitModel,
				        """image-classification""": BeitForImageClassification,
				        """image-segmentation""": BeitForSemanticSegmentation,
				    }
				    if is_torch_available()
				    else {}
				)

				lowerCAmelCase_     			=							False
				lowerCAmelCase_     			=							False
				lowerCAmelCase_     			=							False







				def 						UpperCAmelCase_    (       self     )->		Any:







											'''simple docstring'''






											UpperCamelCase   			=     BeitModelTester(self     )
											UpperCamelCase   			=     ConfigTester(self						,				config_class=A_						,				has_text_modality=A_						,				hidden_size=37     )







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







											'''simple docstring'''






											self.config_tester.run_common_tests()







				@unittest.skip(reason='BEiT does not use inputs_embeds'     )
				def 						UpperCAmelCase_    (       self     )->		Optional[int]:







											'''simple docstring'''






											pass







				@require_torch_multi_gpu
				@unittest.skip(reason='BEiT has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'     )
				def 						UpperCAmelCase_    (       self     )->		Optional[Any]:







											'''simple docstring'''






											pass







				def 						UpperCAmelCase_    (       self     )->		Tuple:







											'''simple docstring'''






											UpperCamelCase			,	UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs_for_common()

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







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







											'''simple docstring'''






											UpperCamelCase			,	UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs_for_common()

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

																		UpperCamelCase   			=     ['pixel_values']
																		self.assertListEqual(arg_names[:1]						,				A_     )







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







											'''simple docstring'''






											UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*A_     )







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







											'''simple docstring'''






											UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*A_     )







				def 						UpperCAmelCase_    (       self     )->		Any:







											'''simple docstring'''






											UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*A_     )







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







											'''simple docstring'''






											UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_semantic_segmentation(*A_     )







				def 						UpperCAmelCase_    (       self     )->		int:







											'''simple docstring'''






											if not self.model_tester.is_training:
																		return

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

											for model_class in self.all_model_classes:
																		# we don't test BeitForMaskedImageModeling
																		if model_class in [*get_values(A_     ), BeitForMaskedImageModeling]:
																									continue

																		UpperCamelCase   			=     model_class(A_     )
																		model.to(A_     )
																		model.train()
																		UpperCamelCase   			=     self._prepare_for_class(A_						,				A_						,				return_labels=A_     )
																		UpperCamelCase   			=     model(**A_     ).loss
																		loss.backward()







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







											'''simple docstring'''






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

											UpperCamelCase   			=     False
											UpperCamelCase   			=     True

											for model_class in self.all_model_classes:
																		# we don't test BeitForMaskedImageModeling
																		if (
																		    model_class in [*get_values(A_     ), BeitForMaskedImageModeling]
																		    or not model_class.supports_gradient_checkpointing
																		):
																									continue

																		UpperCamelCase   			=     model_class(A_     )
																		model.gradient_checkpointing_enable()
																		model.to(A_     )
																		model.train()
																		UpperCamelCase   			=     self._prepare_for_class(A_						,				A_						,				return_labels=A_     )
																		UpperCamelCase   			=     model(**A_     ).loss
																		loss.backward()







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







											'''simple docstring'''






											UpperCamelCase			,	UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs_for_common()

											UpperCamelCase   			=     _config_zero_init(A_     )
											for model_class in self.all_model_classes:
																		UpperCamelCase   			=     model_class(config=A_     )
																		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'''						,				)







				@slow
				def 						UpperCAmelCase_    (       self     )->		Dict:







											'''simple docstring'''






											for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		UpperCamelCase   			=     BeitModel.from_pretrained(A_     )
																		self.assertIsNotNone(A_     )







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





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







				@cached_property
				def 						UpperCAmelCase_    (       self     )->		Optional[int]:







											'''simple docstring'''






											return BeitImageProcessor.from_pretrained('microsoft/beit-base-patch16-224'     ) if is_vision_available() else None







				@slow
				def 						UpperCAmelCase_    (       self     )->		Optional[Any]:







											'''simple docstring'''






											UpperCamelCase   			=     BeitForMaskedImageModeling.from_pretrained('microsoft/beit-base-patch16-224-pt22k'     ).to(A_     )

											UpperCamelCase   			=     self.default_image_processor
											UpperCamelCase   			=     prepare_img()
											UpperCamelCase   			=     image_processor(images=A_						,				return_tensors='pt'     ).pixel_values.to(A_     )

											# prepare bool_masked_pos
											UpperCamelCase   			=     torch.ones((1, 196)						,				dtype=torch.bool     ).to(A_     )

											# forward pass
											with torch.no_grad():
																		UpperCamelCase   			=     model(pixel_values=A_						,				bool_masked_pos=A_     )
											UpperCamelCase   			=     outputs.logits

											# verify the logits
											UpperCamelCase   			=     torch.Size((1, 196, 8192)     )
											self.assertEqual(logits.shape						,				A_     )

											UpperCamelCase   			=     torch.tensor(
											    [[-3.2_437, 0.5_072, -13.9_174], [-3.2_456, 0.4_948, -13.9_401], [-3.2_033, 0.5_121, -13.8_550]]     ).to(A_     )

											self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]						,				A_						,				atol=1e-2     )     )







				@slow
				def 						UpperCAmelCase_    (       self     )->		Union[str, Any]:







											'''simple docstring'''






											UpperCamelCase   			=     BeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224'     ).to(A_     )

											UpperCamelCase   			=     self.default_image_processor
											UpperCamelCase   			=     prepare_img()
											UpperCamelCase   			=     image_processor(images=A_						,				return_tensors='pt'     ).to(A_     )

											# forward pass
											with torch.no_grad():
																		UpperCamelCase   			=     model(**A_     )
											UpperCamelCase   			=     outputs.logits

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

											UpperCamelCase   			=     torch.tensor([-1.2_385, -1.0_987, -1.0_108]     ).to(A_     )

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

											UpperCamelCase   			=     281
											self.assertEqual(logits.argmax(-1     ).item()						,				A_     )







				@slow
				def 						UpperCAmelCase_    (       self     )->		Optional[Any]:







											'''simple docstring'''






											UpperCamelCase   			=     BeitForImageClassification.from_pretrained('microsoft/beit-large-patch16-224-pt22k-ft22k'     ).to(
											    A_     )

											UpperCamelCase   			=     self.default_image_processor
											UpperCamelCase   			=     prepare_img()
											UpperCamelCase   			=     image_processor(images=A_						,				return_tensors='pt'     ).to(A_     )

											# forward pass
											with torch.no_grad():
																		UpperCamelCase   			=     model(**A_     )
											UpperCamelCase   			=     outputs.logits

											# verify the logits
											UpperCamelCase   			=     torch.Size((1, 21841)     )
											self.assertEqual(logits.shape						,				A_     )

											UpperCamelCase   			=     torch.tensor([1.6_881, -0.2_787, 0.5_901]     ).to(A_     )

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

											UpperCamelCase   			=     2396
											self.assertEqual(logits.argmax(-1     ).item()						,				A_     )







				@slow
				def 						UpperCAmelCase_    (       self     )->		Any:







											'''simple docstring'''






											UpperCamelCase   			=     BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640'     )
											UpperCamelCase   			=     model.to(A_     )

											UpperCamelCase   			=     BeitImageProcessor(do_resize=A_						,				size=640						,				do_center_crop=A_     )

											UpperCamelCase   			=     load_dataset('hf-internal-testing/fixtures_ade20k'						,				split='test'     )
											UpperCamelCase   			=     Image.open(ds[0]['file']     )
											UpperCamelCase   			=     image_processor(images=A_						,				return_tensors='pt'     ).to(A_     )

											# forward pass
											with torch.no_grad():
																		UpperCamelCase   			=     model(**A_     )
											UpperCamelCase   			=     outputs.logits

											# verify the logits
											UpperCamelCase   			=     torch.Size((1, 150, 160, 160)     )
											self.assertEqual(logits.shape						,				A_     )

											UpperCamelCase   			=     version.parse(PIL.__version__     ) < version.parse('9.0.0'     )

											if is_pillow_less_than_a:
																		UpperCamelCase   			=     torch.tensor(
																		    [
																		        [[-4.9_225, -2.3_954, -3.0_522], [-2.8_822, -1.0_046, -1.7_561], [-2.9_549, -1.3_228, -2.1_347]],
																		        [[-5.8_168, -3.4_129, -4.0_778], [-3.8_651, -2.2_214, -3.0_277], [-3.8_356, -2.4_643, -3.3_535]],
																		        [[-0.0_078, 3.9_952, 4.0_754], [2.9_856, 4.6_944, 5.0_035], [3.2_413, 4.7_813, 4.9_969]],
																		    ]						,				device=A_						,				)
											else:
																		UpperCamelCase   			=     torch.tensor(
																		    [
																		        [[-4.8_960, -2.3_688, -3.0_355], [-2.8_478, -0.9_836, -1.7_418], [-2.9_449, -1.3_332, -2.1_456]],
																		        [[-5.8_081, -3.4_124, -4.1_006], [-3.8_561, -2.2_081, -3.0_323], [-3.8_365, -2.4_601, -3.3_669]],
																		        [[-0.0_309, 3.9_868, 4.0_540], [2.9_640, 4.6_877, 4.9_976], [3.2_081, 4.7_690, 4.9_942]],
																		    ]						,				device=A_						,				)

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







				@slow
				def 						UpperCAmelCase_    (       self     )->		Tuple:







											'''simple docstring'''






											UpperCamelCase   			=     BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640'     )
											UpperCamelCase   			=     model.to(A_     )

											UpperCamelCase   			=     BeitImageProcessor(do_resize=A_						,				size=640						,				do_center_crop=A_     )

											UpperCamelCase   			=     load_dataset('hf-internal-testing/fixtures_ade20k'						,				split='test'     )
											UpperCamelCase   			=     Image.open(ds[0]['file']     )
											UpperCamelCase   			=     image_processor(images=A_						,				return_tensors='pt'     ).to(A_     )

											# forward pass
											with torch.no_grad():
																		UpperCamelCase   			=     model(**A_     )

											UpperCamelCase   			=     outputs.logits.detach().cpu()

											UpperCamelCase   			=     image_processor.post_process_semantic_segmentation(outputs=A_						,				target_sizes=[(500, 300)]     )
											UpperCamelCase   			=     torch.Size((500, 300)     )
											self.assertEqual(segmentation[0].shape						,				A_     )

											UpperCamelCase   			=     image_processor.post_process_semantic_segmentation(outputs=A_     )
											UpperCamelCase   			=     torch.Size((160, 160)     )
											self.assertEqual(segmentation[0].shape						,				A_     )




