


'''simple docstring'''

import inspect
import unittest

from transformers import BitConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import BitBackbone, BitForImageClassification, BitImageProcessor, BitModel
 from transformers.models.bit.modeling_bit import BIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image
class   __lowerCAmelCase						:





      '''simple docstring'''






      def __init__(self    :    Optional[Any]   ,						UpperCamelCase    :    Union[str, Any]   ,						UpperCamelCase    :    Any=3   ,						UpperCamelCase    :    str=32   ,						UpperCamelCase    :    Tuple=3   ,						UpperCamelCase    :    Union[str, Any]=10   ,						UpperCamelCase    :    int=[8, 16, 32, 64]   ,						UpperCamelCase    :    Dict=[1, 1, 2, 1]   ,						UpperCamelCase    :    int=True   ,						UpperCamelCase    :    List[Any]=True   ,						UpperCamelCase    :    Optional[int]="relu"   ,						UpperCamelCase    :    Optional[Any]=3   ,						UpperCamelCase    :    Tuple=None   ,						UpperCamelCase    :    Optional[int]=["stage2", "stage3", "stage4"]   ,						UpperCamelCase    :    Optional[int]=[2, 3, 4]   ,						UpperCamelCase    :    Tuple=1   ,						):

             '''simple docstring'''







             lowercase__      			= parent
             lowercase__      			= batch_size
             lowercase__      			= image_size
             lowercase__      			= num_channels
             lowercase__      			= embeddings_size
             lowercase__      			= hidden_sizes
             lowercase__      			= depths
             lowercase__      			= is_training
             lowercase__      			= use_labels
             lowercase__      			= hidden_act
             lowercase__      			= num_labels
             lowercase__      			= scope
             lowercase__      			= len(UpperCamelCase )
             lowercase__      			= out_features
             lowercase__      			= out_indices
             lowercase__      			= num_groups



      def       UpperCamelCase__		(self    :    List[Any] ):

             '''simple docstring'''







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

             lowercase__      			= None
             if self.use_labels:
                    lowercase__      			= ids_tensor([self.batch_size]   ,						self.num_labels )

             lowercase__      			= self.get_config()

             return config, pixel_values, labels



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

             '''simple docstring'''







             return BitConfig(
                 num_channels=self.num_channels   ,						embeddings_size=self.embeddings_size   ,						hidden_sizes=self.hidden_sizes   ,						depths=self.depths   ,						hidden_act=self.hidden_act   ,						num_labels=self.num_labels   ,						out_features=self.out_features   ,						out_indices=self.out_indices   ,						num_groups=self.num_groups   ,						)



      def       UpperCamelCase__		(self    :    List[str]   ,						UpperCamelCase    :    Optional[Any]   ,						UpperCamelCase    :    List[Any]   ,						UpperCamelCase    :    str ):

             '''simple docstring'''







             lowercase__      			= BitModel(config=UpperCamelCase )
             model.to(UpperCamelCase )
             model.eval()
             lowercase__      			= model(UpperCamelCase )
             self.parent.assertEqual(
                 result.last_hidden_state.shape   ,						(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)   ,						)



      def       UpperCamelCase__		(self    :    Union[str, Any]   ,						UpperCamelCase    :    Optional[int]   ,						UpperCamelCase    :    str   ,						UpperCamelCase    :    Dict ):

             '''simple docstring'''







             lowercase__      			= self.num_labels
             lowercase__      			= BitForImageClassification(UpperCamelCase )
             model.to(UpperCamelCase )
             model.eval()
             lowercase__      			= model(UpperCamelCase   ,						labels=UpperCamelCase )
             self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.num_labels) )



      def       UpperCamelCase__		(self    :    Any   ,						UpperCamelCase    :    List[Any]   ,						UpperCamelCase    :    Optional[int]   ,						UpperCamelCase    :    str ):

             '''simple docstring'''







             lowercase__      			= BitBackbone(config=UpperCamelCase )
             model.to(UpperCamelCase )
             model.eval()
             lowercase__      			= model(UpperCamelCase )

             # verify feature maps
             self.parent.assertEqual(len(result.feature_maps )   ,						len(config.out_features ) )
             self.parent.assertListEqual(list(result.feature_maps[0].shape )   ,						[self.batch_size, self.hidden_sizes[1], 4, 4] )

             # verify channels
             self.parent.assertEqual(len(model.channels )   ,						len(config.out_features ) )
             self.parent.assertListEqual(model.channels   ,						config.hidden_sizes[1:] )

             # verify backbone works with out_features=None
             lowercase__      			= None
             lowercase__      			= BitBackbone(config=UpperCamelCase )
             model.to(UpperCamelCase )
             model.eval()
             lowercase__      			= model(UpperCamelCase )

             # verify feature maps
             self.parent.assertEqual(len(result.feature_maps )   ,						1 )
             self.parent.assertListEqual(list(result.feature_maps[0].shape )   ,						[self.batch_size, self.hidden_sizes[-1], 1, 1] )

             # verify channels
             self.parent.assertEqual(len(model.channels )   ,						1 )
             self.parent.assertListEqual(model.channels   ,						[config.hidden_sizes[-1]] )






      def       UpperCamelCase__		(self    :    Any ):

             '''simple docstring'''







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






@require_torch
class   __lowerCAmelCase						(lowercase_  ,     lowercase_  ,     unittest.TestCase     ):





      '''simple docstring'''






      lowerCAmelCase__					:     List[Any]  					=							(BitModel, BitForImageClassification, BitBackbone) if is_torch_available() else ()
      lowerCAmelCase__					:     Union[str, Any]  					=							(
          {"""feature-extraction""": BitModel, """image-classification""": BitForImageClassification}
          if is_torch_available()
          else {}
      )

      lowerCAmelCase__					:     Optional[Any]  					=							False
      lowerCAmelCase__					:     str  					=							False
      lowerCAmelCase__					:     Tuple  					=							False
      lowerCAmelCase__					:     int  					=							False
      lowerCAmelCase__					:     str  					=							False



      def       UpperCamelCase__		(self    :    List[Any] ):

             '''simple docstring'''







             lowercase__      			= BitModelTester(self )
             lowercase__      			= ConfigTester(self   ,						config_class=UpperCamelCase   ,						has_text_modality=UpperCamelCase )



      def       UpperCamelCase__		(self    :    str ):

             '''simple docstring'''







             self.create_and_test_config_common_properties()
             self.config_tester.create_and_test_config_to_json_string()
             self.config_tester.create_and_test_config_to_json_file()
             self.config_tester.create_and_test_config_from_and_save_pretrained()
             self.config_tester.create_and_test_config_with_num_labels()
             self.config_tester.check_config_can_be_init_without_params()
             self.config_tester.check_config_arguments_init()



      def       UpperCamelCase__		(self    :    Tuple ):

             '''simple docstring'''







             return



      @unittest.skip(reason='''Bit does not output attentions''' )
      def       UpperCamelCase__		(self    :    Any ):

             '''simple docstring'''







             pass



      @unittest.skip(reason='''Bit does not use inputs_embeds''' )
      def       UpperCamelCase__		(self    :    Any ):

             '''simple docstring'''







             pass



      @unittest.skip(reason='''Bit does not support input and output embeddings''' )
      def       UpperCamelCase__		(self    :    Tuple ):

             '''simple docstring'''







             pass



      def       UpperCamelCase__		(self    :    List[str] ):

             '''simple docstring'''







             lowercase__			,lowercase__      			= self.model_tester.prepare_config_and_inputs_for_common()

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

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



      def       UpperCamelCase__		(self    :    Optional[int] ):

             '''simple docstring'''







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



      def       UpperCamelCase__		(self    :    str ):

             '''simple docstring'''







             lowercase__      			= self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_backbone(*UpperCamelCase )



      def       UpperCamelCase__		(self    :    Optional[int] ):

             '''simple docstring'''







             lowercase__			,lowercase__      			= self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                    lowercase__      			= model_class(config=UpperCamelCase )
                    for name, module in model.named_modules():
                           if isinstance(UpperCamelCase   ,						(nn.BatchNormad, nn.GroupNorm) ):
                                  self.assertTrue(
                                      torch.all(module.weight == 1 )   ,						msg=f"Parameter {name} of model {model_class} seems not properly initialized"   ,						)
                                  self.assertTrue(
                                      torch.all(module.bias == 0 )   ,						msg=f"Parameter {name} of model {model_class} seems not properly initialized"   ,						)



      def       UpperCamelCase__		(self    :    List[str] ):

             '''simple docstring'''







             def check_hidden_states_output(UpperCamelCase    :    List[str]   ,						UpperCamelCase    :    Optional[int]   ,						UpperCamelCase    :    List[str] ):
                    lowercase__      			= model_class(UpperCamelCase )
                    model.to(UpperCamelCase )
                    model.eval()

                    with torch.no_grad():
                           lowercase__      			= model(**self._prepare_for_class(UpperCamelCase   ,						UpperCamelCase ) )

                    lowercase__      			= outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                    lowercase__      			= self.model_tester.num_stages
                    self.assertEqual(len(UpperCamelCase )   ,						expected_num_stages + 1 )

                    # Bit's feature maps are of shape (batch_size, num_channels, height, width)
                    self.assertListEqual(
                        list(hidden_states[0].shape[-2:] )   ,						[self.model_tester.image_size // 4, self.model_tester.image_size // 4]   ,						)

             lowercase__			,lowercase__      			= self.model_tester.prepare_config_and_inputs_for_common()
             lowercase__      			= ['''preactivation''', '''bottleneck''']
             for model_class in self.all_model_classes:
                    for layer_type in layers_type:
                           lowercase__      			= layer_type
                           lowercase__      			= True
                           check_hidden_states_output(UpperCamelCase   ,						UpperCamelCase   ,						UpperCamelCase )

                           # check that output_hidden_states also work using config
                           del inputs_dict["output_hidden_states"]
                           lowercase__      			= True

                           check_hidden_states_output(UpperCamelCase   ,						UpperCamelCase   ,						UpperCamelCase )



      @unittest.skip(reason='''Bit does not use feedforward chunking''' )
      def       UpperCamelCase__		(self    :    Tuple ):

             '''simple docstring'''







             pass



      def       UpperCamelCase__		(self    :    Dict ):

             '''simple docstring'''







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






      @slow
      def       UpperCamelCase__		(self    :    Any ):

             '''simple docstring'''







             for model_name in BIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    lowercase__      			= BitModel.from_pretrained(UpperCamelCase )
                    self.assertIsNotNone(UpperCamelCase )






def _SCREAMING_SNAKE_CASE      ()       ->       List[str]:






       """simple docstring"""







       lowercase__      			= Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''       )
       return image
@require_torch
@require_vision
class   __lowerCAmelCase						(unittest.TestCase     ):





      '''simple docstring'''






      @cached_property
      def       UpperCamelCase__		(self    :    int ):

             '''simple docstring'''







             return (
                 BitImageProcessor.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0] ) if is_vision_available() else None
             )






      @slow
      def       UpperCamelCase__		(self    :    Optional[Any] ):

             '''simple docstring'''







             lowercase__      			= BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0] ).to(UpperCamelCase )

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

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

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

             lowercase__      			= torch.tensor([[-0.65_26, -0.52_63, -1.43_98]] ).to(UpperCamelCase )

             self.assertTrue(torch.allclose(outputs.logits[0, :3]   ,						UpperCamelCase   ,						atol=1E-4 ) )






@require_torch
class   __lowerCAmelCase						(lowercase_  ,     unittest.TestCase     ):





      '''simple docstring'''






      lowerCAmelCase__					:     Dict  					=							(BitBackbone,) if is_torch_available() else ()
      lowerCAmelCase__					:     Optional[int]  					=							BitConfig

      lowerCAmelCase__					:     List[str]  					=							False



      def       UpperCamelCase__		(self    :    List[str] ):

             '''simple docstring'''







             lowercase__      			= BitModelTester(self )








'''simple docstring'''

from ....utils import logging


lowerCamelCase							:   Optional[Any]  						=	logging.get_logger(__name__)
class   __lowerCAmelCase						(lowercase_     ):





      '''simple docstring'''






      def __init__(self    :    Optional[int]   ,						UpperCamelCase    :    Union[str, Any]   ,						UpperCamelCase    :    Optional[Any]=None   ,						UpperCamelCase    :    int=2048 ):

             '''simple docstring'''







             lowercase__      			= config.__dict__
             lowercase__      			= modal_hidden_size
             if num_labels:
                    lowercase__      			= num_labels




