





'''simple docstring'''


import inspect
import unittest

import numpy as np

from transformers import BeitConfig
from transformers.testing_utils import require_flax, require_vision, slow
from transformers.utils import cached_property, is_flax_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor


if is_flax_available():
 import jax

 from transformers import FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling, FlaxBeitModel

if is_vision_available():
 from PIL import Image

 from transformers import BeitImageProcessor
class       a_   (				unittest.TestCase							):





     def __init__( self		:					Dict	,      lowercase		:					Union[str, Any]	,      lowercase		:					Any=100	,      lowercase		:					List[str]=13	,      lowercase		:					int=30	,      lowercase		:					Dict=2	,      lowercase		:					List[str]=3	,      lowercase		:					List[Any]=True	,      lowercase		:					Union[str, Any]=True	,      lowercase		:					Union[str, Any]=32	,      lowercase		:					Tuple=5	,      lowercase		:					List[Any]=4	,      lowercase		:					Union[str, Any]=37	,      lowercase		:					Optional[int]="gelu"	,      lowercase		:					Any=0.1	,      lowercase		:					Optional[Any]=0.1	,      lowercase		:					Optional[Any]=10	,      lowercase		:					Union[str, Any]=0.02	,      lowercase		:					Union[str, Any]=3	,      ):





        """simple docstring"""




        lowercase_			:int    			= parent
        lowercase_			:List[Any]    			= vocab_size
        lowercase_			:Dict    			= batch_size
        lowercase_			:str    			= image_size
        lowercase_			:List[Any]    			= patch_size
        lowercase_			:List[str]    			= num_channels
        lowercase_			:Tuple    			= is_training
        lowercase_			:Dict    			= use_labels
        lowercase_			:Any    			= hidden_size
        lowercase_			:str    			= num_hidden_layers
        lowercase_			:List[Any]    			= num_attention_heads
        lowercase_			:Union[str, Any]    			= intermediate_size
        lowercase_			:List[str]    			= hidden_act
        lowercase_			:Optional[Any]    			= hidden_dropout_prob
        lowercase_			:Optional[Any]    			= attention_probs_dropout_prob
        lowercase_			:Union[str, Any]    			= type_sequence_label_size
        lowercase_			:Union[str, Any]    			= initializer_range

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





     def 							lowercase__      ( self		:					str						):





        """simple docstring"""




        lowercase_			:List[Any]    			= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

        lowercase_			:Tuple    			= None
        if self.use_labels:
           lowercase_			:Optional[Any]    			= ids_tensor([self.batch_size]	,      self.type_sequence_label_size						)

        lowercase_			:Optional[int]    			= 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=lowercase	,      initializer_range=self.initializer_range	,      )

        return config, pixel_values, labels





     def 							lowercase__      ( self		:					Union[str, Any]	,      lowercase		:					str	,      lowercase		:					Optional[Any]	,      lowercase		:					Optional[Any]						):





        """simple docstring"""




        lowercase_			:List[str]    			= FlaxBeitModel(config=lowercase						)
        lowercase_			:Tuple    			= model(lowercase						)
        self.parent.assertEqual(result.last_hidden_state.shape	,      (self.batch_size, self.seq_length, self.hidden_size)						)





     def 							lowercase__      ( self		:					Union[str, Any]	,      lowercase		:					Optional[int]	,      lowercase		:					Union[str, Any]	,      lowercase		:					Dict						):





        """simple docstring"""




        lowercase_			:int    			= FlaxBeitForMaskedImageModeling(config=lowercase						)
        lowercase_			:str    			= model(lowercase						)
        self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.seq_length - 1, self.vocab_size)						)





     def 							lowercase__      ( self		:					List[str]	,      lowercase		:					Union[str, Any]	,      lowercase		:					Union[str, Any]	,      lowercase		:					Dict						):





        """simple docstring"""




        lowercase_			:str    			= self.type_sequence_label_size
        lowercase_			:Optional[Any]    			= FlaxBeitForImageClassification(config=lowercase						)
        lowercase_			:List[str]    			= model(lowercase						)
        self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.type_sequence_label_size)						)

        # test greyscale images
        lowercase_			:Optional[int]    			= 1
        lowercase_			:List[str]    			= FlaxBeitForImageClassification(lowercase						)

        lowercase_			:Dict    			= floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
        lowercase_			:int    			= model(lowercase						)







     def 							lowercase__      ( self		:					Union[str, Any]						):





        """simple docstring"""




        lowercase_			:int    			= self.prepare_config_and_inputs()
        (
            (
            lowercase_
        )			,							(
            lowercase_
        )			,							(
            lowercase_
        )			,							
        )			:List[str]    			= config_and_inputs
        lowercase_			:Optional[Any]    			= {"pixel_values": pixel_values}
        return config, inputs_dict







@require_flax
class       a_   (				_lowerCAmelCase	,						unittest.TestCase							):
     __A      =			(
         (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
     )





     def 							lowercase__      ( self		:					Optional[int]						):





        """simple docstring"""




        lowercase_			:Dict    			= FlaxBeitModelTester(self						)
        lowercase_			:int    			= ConfigTester(self	,      config_class=lowercase	,      has_text_modality=lowercase	,      hidden_size=37						)





     def 							lowercase__      ( self		:					int						):





        """simple docstring"""




        self.config_tester.run_common_tests()





     def 							lowercase__      ( self		:					Tuple						):





        """simple docstring"""




        lowercase_			,							lowercase_			:Optional[int]    			= self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           lowercase_			:Optional[Any]    			= model_class(lowercase						)
           lowercase_			:List[Any]    			= inspect.signature(model.__call__						)
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           lowercase_			:str    			= [*signature.parameters.keys()]

           lowercase_			:List[Any]    			= ["pixel_values"]
           self.assertListEqual(arg_names[:1]	,      lowercase						)





     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        lowercase_			,							lowercase_			:Tuple    			= self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           with self.subTest(model_class.__name__						):
              lowercase_			:Dict    			= self._prepare_for_class(lowercase	,      lowercase						)
              lowercase_			:Optional[int]    			= model_class(lowercase						)

              @jax.jit
              def model_jitted(lowercase		:					Dict	,      **lowercase		:					Any						):
                 return model(pixel_values=lowercase	,      **lowercase						)

              with self.subTest("JIT Enabled"						):
                 lowercase_			:int    			= model_jitted(**lowercase						).to_tuple()

              with self.subTest("JIT Disabled"						):
                 with jax.disable_jit():
                    lowercase_			:Any    			= model_jitted(**lowercase						).to_tuple()

              self.assertEqual(len(lowercase						)	,      len(lowercase						)						)
              for jitted_output, output in zip(lowercase	,      lowercase						):
                 self.assertEqual(jitted_output.shape	,      output.shape						)





     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        lowercase_			:str    			= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase						)





     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        lowercase_			:Any    			= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*lowercase						)





     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:int    			= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*lowercase						)







     @slow
     def 							lowercase__      ( self		:					int						):





        """simple docstring"""




        for model_class_name in self.all_model_classes:
           lowercase_			:str    			= model_class_name.from_pretrained("microsoft/beit-base-patch16-224"						)
           lowercase_			:int    			= model(np.ones((1, 3, 224, 224)						)						)
           self.assertIsNotNone(lowercase						)







def 							UpperCAmelCase_  (						):
   lowercase_			:List[str]    			= Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png" )
   return image


@require_vision
@require_flax
class       a_   (				unittest.TestCase							):





     @cached_property
     def 							lowercase__      ( self		:					Optional[Any]						):





        """simple docstring"""




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





     @slow
     def 							lowercase__      ( self		:					str						):





        """simple docstring"""




        lowercase_			:List[str]    			= FlaxBeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"						)

        lowercase_			:str    			= self.default_image_processor
        lowercase_			:int    			= prepare_img()
        lowercase_			:Any    			= image_processor(images=lowercase	,      return_tensors="np"						).pixel_values

        # prepare bool_masked_pos
        lowercase_			:List[str]    			= np.ones((1, 196)	,      dtype=lowercase						)

        # forward pass
        lowercase_			:List[str]    			= model(pixel_values=lowercase	,      bool_masked_pos=lowercase						)
        lowercase_			:Dict    			= outputs.logits

        # verify the logits
        lowercase_			:Tuple    			= (1, 196, 8_192)
        self.assertEqual(logits.shape	,      lowercase						)

        lowercase_			:Dict    			= np.array(
            [[-3.24_37, 0.50_72, -13.91_74], [-3.24_56, 0.49_48, -13.94_01], [-3.20_33, 0.51_21, -13.85_50]]						)

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





     @slow
     def 							lowercase__      ( self		:					Tuple						):





        """simple docstring"""




        lowercase_			:int    			= FlaxBeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"						)

        lowercase_			:Tuple    			= self.default_image_processor
        lowercase_			:Dict    			= prepare_img()
        lowercase_			:str    			= image_processor(images=lowercase	,      return_tensors="np"						)

        # forward pass
        lowercase_			:Any    			= model(**lowercase						)
        lowercase_			:Tuple    			= outputs.logits

        # verify the logits
        lowercase_			:Dict    			= (1, 1_000)
        self.assertEqual(logits.shape	,      lowercase						)

        lowercase_			:List[str]    			= np.array([-1.23_85, -1.09_87, -1.01_08]						)

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

        lowercase_			:Dict    			= 281
        self.assertEqual(logits.argmax(-1						).item()	,      lowercase						)







     @slow
     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:List[str]    			= FlaxBeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"						)

        lowercase_			:List[Any]    			= self.default_image_processor
        lowercase_			:Union[str, Any]    			= prepare_img()
        lowercase_			:Tuple    			= image_processor(images=lowercase	,      return_tensors="np"						)

        # forward pass
        lowercase_			:List[Any]    			= model(**lowercase						)
        lowercase_			:Dict    			= outputs.logits

        # verify the logits
        lowercase_			:List[str]    			= (1, 21_841)
        self.assertEqual(logits.shape	,      lowercase						)

        lowercase_			:Union[str, Any]    			= np.array([1.68_81, -0.27_87, 0.59_01]						)

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

        lowercase_			:List[Any]    			= 2_396
        self.assertEqual(logits.argmax(-1						).item()	,      lowercase						)









'''simple docstring'''


import unittest

import numpy as np

from transformers import BertConfig, 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():
 from transformers.models.bert.modeling_flax_bert import (
     FlaxBertForMaskedLM,
     FlaxBertForMultipleChoice,
     FlaxBertForNextSentencePrediction,
     FlaxBertForPreTraining,
     FlaxBertForQuestionAnswering,
     FlaxBertForSequenceClassification,
     FlaxBertForTokenClassification,
     FlaxBertModel,
 )
class       a_   (				unittest.TestCase							):





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





        """simple docstring"""




        lowercase_			:List[str]    			= parent
        lowercase_			:Any    			= batch_size
        lowercase_			:Dict    			= seq_length
        lowercase_			:Union[str, Any]    			= is_training
        lowercase_			:Optional[int]    			= use_attention_mask
        lowercase_			:Any    			= use_token_type_ids
        lowercase_			:Union[str, Any]    			= use_labels
        lowercase_			:Dict    			= vocab_size
        lowercase_			:Tuple    			= hidden_size
        lowercase_			:Tuple    			= num_hidden_layers
        lowercase_			:Optional[int]    			= num_attention_heads
        lowercase_			:Optional[Any]    			= intermediate_size
        lowercase_			:str    			= hidden_act
        lowercase_			:Tuple    			= hidden_dropout_prob
        lowercase_			:Optional[Any]    			= attention_probs_dropout_prob
        lowercase_			:Tuple    			= max_position_embeddings
        lowercase_			:Any    			= type_vocab_size
        lowercase_			:int    			= type_sequence_label_size
        lowercase_			:Tuple    			= initializer_range
        lowercase_			:Optional[Any]    			= num_choices





     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:int    			= ids_tensor([self.batch_size, self.seq_length]	,      self.vocab_size						)

        lowercase_			:Union[str, Any]    			= None
        if self.use_attention_mask:
           lowercase_			:Dict    			= random_attention_mask([self.batch_size, self.seq_length]						)

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

        lowercase_			:Optional[Any]    			= BertConfig(
            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=lowercase	,      initializer_range=self.initializer_range	,      )

        return config, input_ids, token_type_ids, attention_mask





     def 							lowercase__      ( self		:					Union[str, Any]						):





        """simple docstring"""




        lowercase_			:int    			= self.prepare_config_and_inputs()
        lowercase_			,							lowercase_			,							lowercase_			,							lowercase_			:Tuple    			= config_and_inputs
        lowercase_			:Tuple    			= {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
        return config, inputs_dict







     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:Any    			= self.prepare_config_and_inputs()
        lowercase_			,							lowercase_			,							lowercase_			,							lowercase_			:Union[str, Any]    			= config_and_inputs

        lowercase_			:Dict    			= True
        lowercase_			:Optional[int]    			= floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
        lowercase_			:str    			= ids_tensor([self.batch_size, self.seq_length]	,      vocab_size=2						)

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







@require_flax
class       a_   (				_lowerCAmelCase	,						unittest.TestCase							):
     __A      =			True

     __A      =			(
         (
             FlaxBertModel,
             FlaxBertForPreTraining,
             FlaxBertForMaskedLM,
             FlaxBertForMultipleChoice,
             FlaxBertForQuestionAnswering,
             FlaxBertForNextSentencePrediction,
             FlaxBertForSequenceClassification,
             FlaxBertForTokenClassification,
             FlaxBertForQuestionAnswering,
         )
         if is_flax_available()
         else ()
     )





     def 							lowercase__      ( self		:					Any						):





        """simple docstring"""




        lowercase_			:Optional[Any]    			= FlaxBertModelTester(self						)







     @slow
     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        lowercase_			:List[str]    			= FlaxBertModel.from_pretrained("bert-base-cased"						)
        lowercase_			:str    			= model(np.ones((1, 1)						)						)
        self.assertIsNotNone(lowercase						)


