




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

    '''simple docstring'''






    def __init__(    self   :			List[str]      ,				UpperCAmelCase__   :			Dict      ,				UpperCAmelCase__   :			Optional[Any]=100      ,				UpperCAmelCase__   :			List[str]=13      ,				UpperCAmelCase__   :			Union[str, Any]=30      ,				UpperCAmelCase__   :			Optional[Any]=2      ,				UpperCAmelCase__   :			List[Any]=3      ,				UpperCAmelCase__   :			Optional[Any]=True      ,				UpperCAmelCase__   :			Optional[Any]=True      ,				UpperCAmelCase__   :			List[str]=32      ,				UpperCAmelCase__   :			List[Any]=4      ,				UpperCAmelCase__   :			Optional[Any]=4      ,				UpperCAmelCase__   :			Tuple=37      ,				UpperCAmelCase__   :			List[str]="gelu"      ,				UpperCAmelCase__   :			List[str]=0.1      ,				UpperCAmelCase__   :			int=0.1      ,				UpperCAmelCase__   :			List[str]=10      ,				UpperCAmelCase__   :			str=0.02      ,				UpperCAmelCase__   :			List[str]=3      ,				UpperCAmelCase__   :			List[str]=None      ,				UpperCAmelCase__   :			Optional[int]=[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   :			Tuple)		->Dict:





           '''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   :			Optional[Any])		->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=UpperCAmelCase__      ,				initializer_range=self.initializer_range      ,				out_indices=self.out_indices      ,				)


    def  SCREAMING_SNAKE_CASE							(    self   :			int      ,				UpperCAmelCase__   :			Optional[int]      ,				UpperCAmelCase__   :			List[str]      ,				UpperCAmelCase__   :			List[str]      ,				UpperCAmelCase__   :			Optional[Any])		->Dict:





           '''simple docstring'''

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


    def  SCREAMING_SNAKE_CASE							(    self   :			str      ,				UpperCAmelCase__   :			Optional[int]      ,				UpperCAmelCase__   :			Optional[Any]      ,				UpperCAmelCase__   :			Optional[Any]      ,				UpperCAmelCase__   :			str)		->Dict:





           '''simple docstring'''

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


    def  SCREAMING_SNAKE_CASE							(    self   :			Tuple      ,				UpperCAmelCase__   :			Optional[Any]      ,				UpperCAmelCase__   :			str      ,				UpperCAmelCase__   :			Tuple      ,				UpperCAmelCase__   :			List[Any])		->str:





           '''simple docstring'''

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

           # test greyscale images
           A__   				=					1
           A__   				=					BeitForImageClassification(UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()

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


    def  SCREAMING_SNAKE_CASE							(    self   :			Any      ,				UpperCAmelCase__   :			str      ,				UpperCAmelCase__   :			Dict      ,				UpperCAmelCase__   :			Optional[Any]      ,				UpperCAmelCase__   :			int)		->Union[str, Any]:





           '''simple docstring'''

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





           '''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        UpperCamelCase_					(				UpperCAmelCase__	,      UpperCAmelCase__	,      unittest.TestCase  ):

    '''simple docstring'''






    UpperCAmelCase__													=   (
        (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
        if is_torch_available()
        else ()
    )
    UpperCAmelCase__													=   (
        {
            '''feature-extraction''': BeitModel,
            '''image-classification''': BeitForImageClassification,
            '''image-segmentation''': BeitForSemanticSegmentation,
        }
        if is_torch_available()
        else {}
    )

    UpperCAmelCase__													=   False
    UpperCAmelCase__													=   False
    UpperCAmelCase__													=   False


    def  SCREAMING_SNAKE_CASE							(    self   :			List[Any])		->List[str]:





           '''simple docstring'''

           A__   				=					BeitModelTester(self)
           A__   				=					ConfigTester(self      ,				config_class=UpperCAmelCase__      ,				has_text_modality=UpperCAmelCase__      ,				hidden_size=37)


    def  SCREAMING_SNAKE_CASE							(    self   :			Dict)		->int:





           '''simple docstring'''

           self.config_tester.run_common_tests()


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





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





           '''simple docstring'''

           pass


    def  SCREAMING_SNAKE_CASE							(    self   :			List[str])		->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(UpperCAmelCase__)
                  self.assertIsInstance(model.get_input_embeddings()      ,				(nn.Module))
                  A__   				=					model.get_output_embeddings()
                  self.assertTrue(x is None or isinstance(UpperCAmelCase__      ,				nn.Linear))


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





           '''simple docstring'''

           A__				,							A__   				=					self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  A__   				=					model_class(UpperCAmelCase__)
                  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]      ,				UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			Tuple)		->Dict:





           '''simple docstring'''

           A__   				=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			Optional[Any])		->Tuple:





           '''simple docstring'''

           A__   				=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			Optional[int])		->Dict:





           '''simple docstring'''

           A__   				=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			Optional[Any])		->Optional[int]:





           '''simple docstring'''

           A__   				=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_semantic_segmentation(*UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			int)		->Any:





           '''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:
                  # we don't test BeitForMaskedImageModeling
                  if model_class in [*get_values(UpperCAmelCase__), BeitForMaskedImageModeling]:
                         continue

                  A__   				=					model_class(UpperCAmelCase__)
                  model.to(UpperCAmelCase__)
                  model.train()
                  A__   				=					self._prepare_for_class(UpperCAmelCase__      ,				UpperCAmelCase__      ,				return_labels=UpperCAmelCase__)
                  A__   				=					model(**UpperCAmelCase__).loss
                  loss.backward()


    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->List[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:
                  # we don't test BeitForMaskedImageModeling
                  if (
                      model_class in [*get_values(UpperCAmelCase__), BeitForMaskedImageModeling]
                      or not model_class.supports_gradient_checkpointing
                  ):
                         continue

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


    def  SCREAMING_SNAKE_CASE							(    self   :			Dict)		->Dict:





           '''simple docstring'''

           A__				,							A__   				=					self.model_tester.prepare_config_and_inputs_for_common()

           A__   				=					_config_zero_init(UpperCAmelCase__)
           for model_class in self.all_model_classes:
                  A__   				=					model_class(config=UpperCAmelCase__)
                  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  SCREAMING_SNAKE_CASE							(    self   :			int)		->List[Any]:





           '''simple docstring'''

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







def SCREAMING_SNAKE_CASE       ( )     -> str:



       """simple docstring"""

       A__   				=					Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''		)
       return image


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

    '''simple docstring'''






    @cached_property
    def  SCREAMING_SNAKE_CASE							(    self   :			Dict)		->Optional[Any]:





           '''simple docstring'''

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


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





           '''simple docstring'''

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

           A__   				=					self.default_image_processor
           A__   				=					prepare_img()
           A__   				=					image_processor(images=UpperCAmelCase__      ,				return_tensors='''pt''').pixel_values.to(UpperCAmelCase__)

           # prepare bool_masked_pos
           A__   				=					torch.ones((1, 196)      ,				dtype=torch.bool).to(UpperCAmelCase__)

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

           # verify the logits
           A__   				=					torch.Size((1, 196, 8_192))
           self.assertEqual(logits.shape      ,				UpperCAmelCase__)

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

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


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





           '''simple docstring'''

           A__   				=					BeitForImageClassification.from_pretrained('''microsoft/beit-base-patch16-224''').to(UpperCAmelCase__)

           A__   				=					self.default_image_processor
           A__   				=					prepare_img()
           A__   				=					image_processor(images=UpperCAmelCase__      ,				return_tensors='''pt''').to(UpperCAmelCase__)

           # forward pass
           with torch.no_grad():
                  A__   				=					model(**UpperCAmelCase__)
           A__   				=					outputs.logits

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

           A__   				=					torch.tensor([-1.2385, -1.0987, -1.0108]).to(UpperCAmelCase__)

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

           A__   				=					281
           self.assertEqual(logits.argmax(-1).item()      ,				UpperCAmelCase__)


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





           '''simple docstring'''

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

           A__   				=					self.default_image_processor
           A__   				=					prepare_img()
           A__   				=					image_processor(images=UpperCAmelCase__      ,				return_tensors='''pt''').to(UpperCAmelCase__)

           # forward pass
           with torch.no_grad():
                  A__   				=					model(**UpperCAmelCase__)
           A__   				=					outputs.logits

           # verify the logits
           A__   				=					torch.Size((1, 21_841))
           self.assertEqual(logits.shape      ,				UpperCAmelCase__)

           A__   				=					torch.tensor([1.6881, -0.2787, 0.5901]).to(UpperCAmelCase__)

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

           A__   				=					2_396
           self.assertEqual(logits.argmax(-1).item()      ,				UpperCAmelCase__)


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





           '''simple docstring'''

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

           A__   				=					BeitImageProcessor(do_resize=UpperCAmelCase__      ,				size=640      ,				do_center_crop=UpperCAmelCase__)

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

           # forward pass
           with torch.no_grad():
                  A__   				=					model(**UpperCAmelCase__)
           A__   				=					outputs.logits

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

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

           if is_pillow_less_than_a:
                  A__   				=					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=UpperCAmelCase__      ,				)
           else:
                  A__   				=					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=UpperCAmelCase__      ,				)

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





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





           '''simple docstring'''

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

           A__   				=					BeitImageProcessor(do_resize=UpperCAmelCase__      ,				size=640      ,				do_center_crop=UpperCAmelCase__)

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

           # forward pass
           with torch.no_grad():
                  A__   				=					model(**UpperCAmelCase__)

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

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

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








import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
  from transformers import AutoModelForSeqaSeqLM, AutoTokenizer


@require_torch
@require_sentencepiece
@require_tokenizers
class        UpperCamelCase_					(				unittest.TestCase  ):

    '''simple docstring'''






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





           '''simple docstring'''

           A__   				=					AutoModelForSeqaSeqLM.from_pretrained('''google/mt5-small'''      ,				return_dict=UpperCAmelCase__).to(UpperCAmelCase__)
           A__   				=					AutoTokenizer.from_pretrained('''google/mt5-small''')

           A__   				=					tokenizer('''Hello there'''      ,				return_tensors='''pt''').input_ids
           A__   				=					tokenizer('''Hi I am'''      ,				return_tensors='''pt''').input_ids

           A__   				=					model(input_ids.to(UpperCAmelCase__)      ,				labels=labels.to(UpperCAmelCase__)).loss
           A__   				=					-(labels.shape[-1] * loss.item())

           A__   				=					-84.9127
           self.assertTrue(abs(mtf_score - EXPECTED_SCORE) < 1e-4)


