


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



		"""simple docstring"""

		def __init__( self		,      A		,      A=1_00		,      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]		,      )  -> Union[str, Any]:







					'''simple docstring'''







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

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







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







					'''simple docstring'''







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

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

					lowerCamelCase    		=  self.get_config()

					return config, pixel_values, labels, pixel_labels







		def 					__A   ( self	)  -> int:







					'''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 					__A   ( self		,      A		,      A		,      A		,      A	)  -> Dict:







					'''simple docstring'''







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







		def 					__A   ( self		,      A		,      A		,      A		,      A	)  -> str:







					'''simple docstring'''







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







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







					'''simple docstring'''







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

					# test greyscale images
					lowerCamelCase    		=  1
					lowerCamelCase    		=  BeitForImageClassification(A	)
					model.to(A	)
					model.eval()

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







		def 					__A   ( self		,      A		,      A		,      A		,      A	)  -> List[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  self.num_labels
					lowerCamelCase    		=  BeitForSemanticSegmentation(A	)
					model.to(A	)
					model.eval()
					lowerCamelCase    		=  model(A	)
					self.parent.assertEqual(
					    result.logits.shape		,      (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)	)
					lowerCamelCase    		=  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 					__A   ( self	)  -> Optional[int]:







					'''simple docstring'''







					lowerCamelCase    		=  self.prepare_config_and_inputs()
					lowerCamelCase		,		lowerCamelCase		,		lowerCamelCase		,		lowerCamelCase    		=  config_and_inputs
					lowerCamelCase    		=  {"""pixel_values""": pixel_values}
					return config, inputs_dict







@require_torch
class        __lowercase       (							a_			,							a_			,							unittest.TestCase  ):



		"""simple docstring"""

		UpperCamelCase     :					List[str]			       =		(
		    (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
		    if is_torch_available()
		    else ()
		)
		UpperCamelCase     :					List[str]			       =		(
		    {
		        "feature-extraction": BeitModel,
		        "image-classification": BeitForImageClassification,
		        "image-segmentation": BeitForSemanticSegmentation,
		    }
		    if is_torch_available()
		    else {}
		)

		UpperCamelCase     :					int			       =		False
		UpperCamelCase     :					Optional[int]			       =		False
		UpperCamelCase     :					int			       =		False







		def 					__A   ( self	)  -> int:







					'''simple docstring'''







					lowerCamelCase    		=  BeitModelTester(self	)
					lowerCamelCase    		=  ConfigTester(self		,      config_class=A		,      has_text_modality=A		,      hidden_size=37	)







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







					'''simple docstring'''







					self.config_tester.run_common_tests()







		@unittest.skip(reason="""BEiT does not use inputs_embeds"""	)
		def 					__A   ( self	)  -> List[str]:







					'''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 					__A   ( self	)  -> List[Any]:







					'''simple docstring'''







					pass







		def 					__A   ( self	)  -> Optional[int]:







					'''simple docstring'''







					lowerCamelCase		,		lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs_for_common()

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







		def 					__A   ( self	)  -> Optional[int]:







					'''simple docstring'''







					lowerCamelCase		,		lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs_for_common()

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

								lowerCamelCase    		=  ["""pixel_values"""]
								self.assertListEqual(arg_names[:1]		,      A	)







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







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*A	)







		def 					__A   ( self	)  -> Any:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_lm(*A	)







		def 					__A   ( self	)  -> Optional[int]:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_image_classification(*A	)







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







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_semantic_segmentation(*A	)







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







					'''simple docstring'''







					if not self.model_tester.is_training:
								return

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

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

								lowerCamelCase    		=  model_class(A	)
								model.to(A	)
								model.train()
								lowerCamelCase    		=  self._prepare_for_class(A		,      A		,      return_labels=A	)
								lowerCamelCase    		=  model(**A	).loss
								loss.backward()







		def 					__A   ( self	)  -> Tuple:







					'''simple docstring'''







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

					lowerCamelCase    		=  False
					lowerCamelCase    		=  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

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







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







					'''simple docstring'''







					lowerCamelCase		,		lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs_for_common()

					lowerCamelCase    		=  _config_zero_init(A	)
					for model_class in self.all_model_classes:
								lowerCamelCase    		=  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 					__A   ( self	)  -> List[Any]:







					'''simple docstring'''







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






def 	__lowerCamelCase						(							):





			'''simple docstring'''




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







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



		"""simple docstring"""

		@cached_property
		def 					__A   ( self	)  -> Tuple:







					'''simple docstring'''







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







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







					'''simple docstring'''







					lowerCamelCase    		=  BeitForMaskedImageModeling.from_pretrained("""microsoft/beit-base-patch16-224-pt22k"""	).to(A	)

					lowerCamelCase    		=  self.default_image_processor
					lowerCamelCase    		=  prepare_img()
					lowerCamelCase    		=  image_processor(images=A		,      return_tensors="""pt"""	).pixel_values.to(A	)

					# prepare bool_masked_pos
					lowerCamelCase    		=  torch.ones((1, 1_96)		,      dtype=torch.bool	).to(A	)

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

					# verify the logits
					lowerCamelCase    		=  torch.Size((1, 1_96, 81_92)	)
					self.assertEqual(logits.shape		,      A	)

					lowerCamelCase    		=  torch.tensor(
					    [[-3.2437, 0.5072, -13.9174], [-3.2456, 0.4948, -13.9401], [-3.2033, 0.5121, -13.8550]]	).to(A	)

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







		@slow
		def 					__A   ( self	)  -> Optional[int]:







					'''simple docstring'''







					lowerCamelCase    		=  BeitForImageClassification.from_pretrained("""microsoft/beit-base-patch16-224"""	).to(A	)

					lowerCamelCase    		=  self.default_image_processor
					lowerCamelCase    		=  prepare_img()
					lowerCamelCase    		=  image_processor(images=A		,      return_tensors="""pt"""	).to(A	)

					# forward pass
					with torch.no_grad():
								lowerCamelCase    		=  model(**A	)
					lowerCamelCase    		=  outputs.logits

					# verify the logits
					lowerCamelCase    		=  torch.Size((1, 10_00)	)
					self.assertEqual(logits.shape		,      A	)

					lowerCamelCase    		=  torch.tensor([-1.2385, -1.0987, -1.0108]	).to(A	)

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

					lowerCamelCase    		=  2_81
					self.assertEqual(logits.argmax(-1	).item()		,      A	)







		@slow
		def 					__A   ( self	)  -> List[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  BeitForImageClassification.from_pretrained("""microsoft/beit-large-patch16-224-pt22k-ft22k"""	).to(
					    A	)

					lowerCamelCase    		=  self.default_image_processor
					lowerCamelCase    		=  prepare_img()
					lowerCamelCase    		=  image_processor(images=A		,      return_tensors="""pt"""	).to(A	)

					# forward pass
					with torch.no_grad():
								lowerCamelCase    		=  model(**A	)
					lowerCamelCase    		=  outputs.logits

					# verify the logits
					lowerCamelCase    		=  torch.Size((1, 2_18_41)	)
					self.assertEqual(logits.shape		,      A	)

					lowerCamelCase    		=  torch.tensor([1.6881, -0.2787, 0.5901]	).to(A	)

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

					lowerCamelCase    		=  23_96
					self.assertEqual(logits.argmax(-1	).item()		,      A	)







		@slow
		def 					__A   ( self	)  -> List[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  BeitForSemanticSegmentation.from_pretrained("""microsoft/beit-base-finetuned-ade-640-640"""	)
					lowerCamelCase    		=  model.to(A	)

					lowerCamelCase    		=  BeitImageProcessor(do_resize=A		,      size=6_40		,      do_center_crop=A	)

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

					# forward pass
					with torch.no_grad():
								lowerCamelCase    		=  model(**A	)
					lowerCamelCase    		=  outputs.logits

					# verify the logits
					lowerCamelCase    		=  torch.Size((1, 1_50, 1_60, 1_60)	)
					self.assertEqual(logits.shape		,      A	)

					lowerCamelCase    		=  version.parse(PIL.__version__	) < version.parse("""9.0.0"""	)

					if is_pillow_less_than_a:
								lowerCamelCase    		=  torch.tensor(
								    [
								        [[-4.9225, -2.3954, -3.0522], [-2.8822, -1.0046, -1.7561], [-2.9549, -1.3228, -2.1347]],
								        [[-5.8168, -3.4129, -4.0778], [-3.8651, -2.2214, -3.0277], [-3.8356, -2.4643, -3.3535]],
								        [[-0.0078, 3.9952, 4.0754], [2.9856, 4.6944, 5.0035], [3.2413, 4.7813, 4.9969]],
								    ]		,      device=A		,      )
					else:
								lowerCamelCase    		=  torch.tensor(
								    [
								        [[-4.8960, -2.3688, -3.0355], [-2.8478, -0.9836, -1.7418], [-2.9449, -1.3332, -2.1456]],
								        [[-5.8081, -3.4124, -4.1006], [-3.8561, -2.2081, -3.0323], [-3.8365, -2.4601, -3.3669]],
								        [[-0.0309, 3.9868, 4.0540], [2.9640, 4.6877, 4.9976], [3.2081, 4.7690, 4.9942]],
								    ]		,      device=A		,      )

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







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







					'''simple docstring'''







					lowerCamelCase    		=  BeitForSemanticSegmentation.from_pretrained("""microsoft/beit-base-finetuned-ade-640-640"""	)
					lowerCamelCase    		=  model.to(A	)

					lowerCamelCase    		=  BeitImageProcessor(do_resize=A		,      size=6_40		,      do_center_crop=A	)

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

					# forward pass
					with torch.no_grad():
								lowerCamelCase    		=  model(**A	)

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

					lowerCamelCase    		=  image_processor.post_process_semantic_segmentation(outputs=A		,      target_sizes=[(5_00, 3_00)]	)
					lowerCamelCase    		=  torch.Size((5_00, 3_00)	)
					self.assertEqual(segmentation[0].shape		,      A	)

					lowerCamelCase    		=  image_processor.post_process_semantic_segmentation(outputs=A	)
					lowerCamelCase    		=  torch.Size((1_60, 1_60)	)
					self.assertEqual(segmentation[0].shape		,      A	)





# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from .config import config_command_parser
from .config_args import default_config_file, load_config_from_file  # noqa: F401
from .default import default_command_parser
from .update import update_command_parser






def 	__lowerCamelCase						(							lowerCamelCase__		:						Any=None      ):





			'''simple docstring'''




			lowerCamelCase    		=  argparse.ArgumentParser(add_help=lowerCamelCase__  ,							allow_abbrev=lowerCamelCase__      )
			# The main config parser
			lowerCamelCase    		=  config_command_parser(lowerCamelCase__      )
			# The subparser to add commands to
			lowerCamelCase    		=  config_parser.add_subparsers(title="""subcommands"""  ,							dest="""subcommand"""      )

			# Then add other parsers with the parent parser
			default_command_parser(lowerCamelCase__  ,							parents=[parent_parser]      )
			update_command_parser(lowerCamelCase__  ,							parents=[parent_parser]      )

			return config_parser






def 	__lowerCamelCase						(							):





			'''simple docstring'''




			lowerCamelCase    		=  get_config_parser()
			lowerCamelCase    		=  config_parser.parse_args()

			if not hasattr(lowerCamelCase__  ,							"""func"""      ):
						config_parser.print_help()
						exit(1      )

			# Run
			args.func(lowerCamelCase__      )


if __name__ == "__main__":
		main()

