






import unittest

import numpy as np

from transformers import RobertaPreLayerNormConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
     import jax.numpy as jnp

     from transformers.models.roberta_prelayernorm.modeling_flax_roberta_prelayernorm import (
         FlaxRobertaPreLayerNormForCausalLM,
         FlaxRobertaPreLayerNormForMaskedLM,
         FlaxRobertaPreLayerNormForMultipleChoice,
         FlaxRobertaPreLayerNormForQuestionAnswering,
         FlaxRobertaPreLayerNormForSequenceClassification,
         FlaxRobertaPreLayerNormForTokenClassification,
         FlaxRobertaPreLayerNormModel,
     )






class 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




      def __init__( self :		List[Any]					,   __UpperCAmelCase :		Optional[int]					,   __UpperCAmelCase :		Dict=13					,   __UpperCAmelCase :		List[str]=7					,   __UpperCAmelCase :		Tuple=True					,   __UpperCAmelCase :		List[str]=True					,   __UpperCAmelCase :		Optional[Any]=True					,   __UpperCAmelCase :		str=True					,   __UpperCAmelCase :		int=99					,   __UpperCAmelCase :		Union[str, Any]=32					,   __UpperCAmelCase :		Tuple=5					,   __UpperCAmelCase :		Tuple=4					,   __UpperCAmelCase :		Optional[int]=37					,   __UpperCAmelCase :		Optional[Any]="gelu"					,   __UpperCAmelCase :		Any=0.1					,   __UpperCAmelCase :		Optional[int]=0.1					,   __UpperCAmelCase :		Optional[Any]=512					,   __UpperCAmelCase :		Tuple=16					,   __UpperCAmelCase :		List[str]=2					,   __UpperCAmelCase :		int=0.02					,   __UpperCAmelCase :		List[Any]=4					,   )       ->List[Any]:




       """simple docstring"""


       a						    = parent
       a						    = batch_size
       a						    = seq_length
       a						    = is_training
       a						    = use_attention_mask
       a						    = use_token_type_ids
       a						    = use_labels
       a						    = vocab_size
       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						    = max_position_embeddings
       a						    = type_vocab_size
       a						    = type_sequence_label_size
       a						    = initializer_range
       a						    = num_choices







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




       """simple docstring"""


       a						    = ids_tensor([self.batch_size, self.seq_length]					,   self.vocab_size    )

       a						    = None
       if self.use_attention_mask:
        a						    = random_attention_mask([self.batch_size, self.seq_length]    )

       a						    = None
       if self.use_token_type_ids:
        a						    = ids_tensor([self.batch_size, self.seq_length]					,   self.type_vocab_size    )

       a						    = RobertaPreLayerNormConfig(
           vocab_size=self.vocab_size					,   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					,   max_position_embeddings=self.max_position_embeddings					,   type_vocab_size=self.type_vocab_size					,   is_decoder=__UpperCAmelCase					,   initializer_range=self.initializer_range					,   )

       return config, input_ids, token_type_ids, attention_mask







      def  __lowerCAmelCase   ( self :		List[str]    )       ->Tuple:




       """simple docstring"""


       a						    = self.prepare_config_and_inputs()
       a    ,	a    ,	a    ,	a						    = config_and_inputs
       a						    = {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
       return config, inputs_dict







      def  __lowerCAmelCase   ( self :		List[Any]    )       ->Optional[Any]:




       """simple docstring"""


       a						    = self.prepare_config_and_inputs()
       a    ,	a    ,	a    ,	a						    = config_and_inputs

       a						    = True
       a						    = floats_tensor([self.batch_size, self.seq_length, self.hidden_size]    )
       a						    = ids_tensor([self.batch_size, self.seq_length]					,   vocab_size=2    )

       return (
           config,
           input_ids,
           token_type_ids,
           encoder_hidden_states,
           encoder_attention_mask,
       )






@require_flax
# Copied from tests.models.roberta.test_modelling_flax_roberta.FlaxRobertaPreLayerNormModelTest with ROBERTA->ROBERTA_PRELAYERNORM,Roberta->RobertaPreLayerNorm,roberta-base->andreasmadsen/efficient_mlm_m0.40
class 						lowercase_   (      lowercase					,    unittest.TestCase  ):






      '''simple docstring'''




      __snake_case	    =    True

      __snake_case	    =    (
          (
              FlaxRobertaPreLayerNormModel,
              FlaxRobertaPreLayerNormForCausalLM,
              FlaxRobertaPreLayerNormForMaskedLM,
              FlaxRobertaPreLayerNormForSequenceClassification,
              FlaxRobertaPreLayerNormForTokenClassification,
              FlaxRobertaPreLayerNormForMultipleChoice,
              FlaxRobertaPreLayerNormForQuestionAnswering,
          )
          if is_flax_available()
          else ()
      )







      def  __lowerCAmelCase   ( self :		Any    )       ->str:




       """simple docstring"""


       a						    = FlaxRobertaPreLayerNormModelTester(self    )







      @slow
      def  __lowerCAmelCase   ( self :		int    )       ->str:




       """simple docstring"""


       for model_class_name in self.all_model_classes:
        a						    = model_class_name.from_pretrained('''andreasmadsen/efficient_mlm_m0.40'''					,   from_pt=__UpperCAmelCase    )
        a						    = model(np.ones((1, 1)    )    )
        self.assertIsNotNone(__UpperCAmelCase    )






@require_flax
class 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




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




       """simple docstring"""


       a						    = FlaxRobertaPreLayerNormForMaskedLM.from_pretrained('''andreasmadsen/efficient_mlm_m0.40'''					,   from_pt=__UpperCAmelCase    )

       a						    = np.array([[0, 31_414, 232, 328, 740, 1_140, 12_695, 69, 46_078, 1_588, 2]]					,   dtype=jnp.intaa    )
       a						    = model(__UpperCAmelCase    )[0]
       a						    = [1, 11, 50_265]
       self.assertEqual(list(output.shape    )					,   __UpperCAmelCase    )
       # compare the actual values for a slice.
       a						    = np.array(
           [[[40.4880, 18.0199, -5.2367], [-1.8877, -4.0885, 10.7085], [-2.2613, -5.6110, 7.2665]]]					,   dtype=np.floataa    )
       self.assertTrue(np.allclose(output[:, :3, :3]					,   __UpperCAmelCase					,   atol=1e-4    )    )







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




       """simple docstring"""


       a						    = FlaxRobertaPreLayerNormModel.from_pretrained('''andreasmadsen/efficient_mlm_m0.40'''					,   from_pt=__UpperCAmelCase    )

       a						    = np.array([[0, 31_414, 232, 328, 740, 1_140, 12_695, 69, 46_078, 1_588, 2]]					,   dtype=jnp.intaa    )
       a						    = model(__UpperCAmelCase    )[0]
       # compare the actual values for a slice.
       a						    = np.array(
           [[[0.0208, -0.0356, 0.0237], [-0.1569, -0.0411, -0.2626], [0.1879, 0.0125, -0.0089]]]					,   dtype=np.floataa    )
       self.assertTrue(np.allclose(output[:, :3, :3]					,   __UpperCAmelCase					,   atol=1e-4    )    )










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




       """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  __lowerCAmelCase   ( self :		Dict    )       ->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  __lowerCAmelCase   ( self :		List[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  __lowerCAmelCase   ( self :		Tuple					,   __UpperCAmelCase :		Optional[int]					,   __UpperCAmelCase :		Dict					,   __UpperCAmelCase :		List[str]					,   __UpperCAmelCase :		Tuple    )       ->List[Any]:




       """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  __lowerCAmelCase   ( self :		Dict					,   __UpperCAmelCase :		List[Any]					,   __UpperCAmelCase :		Any					,   __UpperCAmelCase :		List[Any]					,   __UpperCAmelCase :		List[str]    )       ->List[str]:




       """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  __lowerCAmelCase   ( self :		Any					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		List[Any]					,   __UpperCAmelCase :		Tuple					,   __UpperCAmelCase :		str    )       ->Optional[int]:




       """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  __lowerCAmelCase   ( self :		int					,   __UpperCAmelCase :		Dict					,   __UpperCAmelCase :		int					,   __UpperCAmelCase :		Tuple					,   __UpperCAmelCase :		Tuple    )       ->int:




       """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  __lowerCAmelCase   ( self :		int    )       ->List[str]:




       """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 						lowercase_   (      lowercase					,    lowercase					,    unittest.TestCase  ):






      '''simple docstring'''




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

      __snake_case	    =    False
      __snake_case	    =    False
      __snake_case	    =    False







      def  __lowerCAmelCase   ( self :		str    )       ->Tuple:




       """simple docstring"""


       a						    = BeitModelTester(self    )
       a						    = ConfigTester(self					,   config_class=__UpperCAmelCase					,   has_text_modality=__UpperCAmelCase					,   hidden_size=37    )







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




       """simple docstring"""


       self.config_tester.run_common_tests()







      @unittest.skip(reason='''BEiT does not use inputs_embeds'''    )
      def  __lowerCAmelCase   ( self :		Any    )       ->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  __lowerCAmelCase   ( self :		Optional[Any]    )       ->Union[str, Any]:




       """simple docstring"""


       pass







      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->str:




       """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  __lowerCAmelCase   ( self :		Optional[int]    )       ->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(__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  __lowerCAmelCase   ( self :		int    )       ->Dict:




       """simple docstring"""


       a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->List[Any]:




       """simple docstring"""


       a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		int    )       ->str:




       """simple docstring"""


       a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Tuple    )       ->List[str]:




       """simple docstring"""


       a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_semantic_segmentation(*__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		int    )       ->Tuple:




       """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  __lowerCAmelCase   ( 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  __lowerCAmelCase   ( self :		Optional[Any]    )       ->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  __lowerCAmelCase   ( self :		List[str]    )       ->Optional[int]:




       """simple docstring"""


       for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        a						    = BeitModel.from_pretrained(__UpperCAmelCase    )
        self.assertIsNotNone(__UpperCAmelCase    )
def 		_a       ( )				->     int:
 a						    = Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''					)
 return image






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






      '''simple docstring'''




      @cached_property
      def  __lowerCAmelCase   ( self :		List[str]    )       ->Union[str, Any]:




       """simple docstring"""


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







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




       """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  __lowerCAmelCase   ( self :		int    )       ->Optional[int]:




       """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  __lowerCAmelCase   ( self :		int    )       ->str:




       """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  __lowerCAmelCase   ( self :		Tuple    )       ->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

       # 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  __lowerCAmelCase   ( self :		str    )       ->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    )


