
'''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 UpperCAmelCase__	:





  """simple docstring"""

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


     '''simple docstring'''






     _a				:	Optional[Any]			=   parent
     _a				:	Any			=   batch_size
     _a				:	Any			=   image_size
     _a				:	Union[str, Any]			=   num_channels
     _a				:	Optional[int]			=   embeddings_size
     _a				:	str			=   hidden_sizes
     _a				:	Optional[Any]			=   depths
     _a				:	Optional[Any]			=   is_training
     _a				:	int			=   use_labels
     _a				:	Union[str, Any]			=   hidden_act
     _a				:	Optional[int]			=   num_labels
     _a				:	Optional[int]			=   scope
     _a				:	Tuple			=   len(_a       )
     _a				:	Union[str, Any]			=   out_features
     _a				:	List[str]			=   out_indices
     _a				:	Any			=   num_groups


  def   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     _a				:	Optional[Any]			=   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

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

     _a				:	Optional[int]			=   self.get_config()

     return config, pixel_values, labels


  def   __lowercase		(		self      :     List[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   __lowercase		(		self      :     Optional[int]							,_a      :     List[str]							,_a      :     Tuple							,_a      :     Union[str, Any]       ):


     '''simple docstring'''






     _a				:	List[Any]			=   BitModel(config=_a       )
     model.to(_a       )
     model.eval()
     _a				:	int			=   model(_a       )
     self.parent.assertEqual(
         result.last_hidden_state.shape							,(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)							,)


  def   __lowercase		(		self      :     List[str]							,_a      :     Dict							,_a      :     str							,_a      :     List[Any]       ):


     '''simple docstring'''






     _a				:	Union[str, Any]			=   self.num_labels
     _a				:	Union[str, Any]			=   BitForImageClassification(_a       )
     model.to(_a       )
     model.eval()
     _a				:	List[Any]			=   model(_a							,labels=_a       )
     self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.num_labels)       )


  def   __lowercase		(		self      :     str							,_a      :     str							,_a      :     Any							,_a      :     Union[str, Any]       ):


     '''simple docstring'''






     _a				:	List[Any]			=   BitBackbone(config=_a       )
     model.to(_a       )
     model.eval()
     _a				:	List[str]			=   model(_a       )

     # 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
     _a				:	Union[str, Any]			=   None
     _a				:	List[Any]			=   BitBackbone(config=_a       )
     model.to(_a       )
     model.eval()
     _a				:	Optional[int]			=   model(_a       )

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


     '''simple docstring'''






     _a				:	int			=   self.prepare_config_and_inputs()
     _a				:	Optional[int]			=   config_and_inputs
     _a				:	str			=   {'pixel_values': pixel_values}
     return config, inputs_dict







@require_torch
class UpperCAmelCase__	( __UpperCAmelCase    ,				__UpperCAmelCase    ,				unittest.TestCase	):





  """simple docstring"""

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

  __UpperCAmelCase				:				Optional[Any]   =					False
  __UpperCAmelCase				:				int   =					False
  __UpperCAmelCase				:				Optional[int]   =					False
  __UpperCAmelCase				:				str   =					False
  __UpperCAmelCase				:				Optional[Any]   =					False


  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     _a				:	List[str]			=   BitModelTester(self       )
     _a				:	int			=   ConfigTester(self							,config_class=_a							,has_text_modality=_a       )


  def   __lowercase		(		self      :     List[Any]       ):


     '''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   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     return


  @unittest.skip(reason='Bit does not output attentions'       )
  def   __lowercase		(		self      :     Optional[int]       ):


     '''simple docstring'''






     pass


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


     '''simple docstring'''






     pass


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


     '''simple docstring'''






     pass


  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     _a				:	Tuple			=   self.model_tester.prepare_config_and_inputs_for_common()

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

        _a				:	Any			=   ['pixel_values']
        self.assertListEqual(arg_names[:1]							,_a       )


  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     _a				:	List[str]			=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*_a       )


  def   __lowercase		(		self      :     int       ):


     '''simple docstring'''






     _a				:	List[str]			=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_backbone(*_a       )


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


     '''simple docstring'''






     _a				:	List[str]			=   self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        _a				:	Union[str, Any]			=   model_class(config=_a       )
        for name, module in model.named_modules():
           if isinstance(_a							,(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   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     def check_hidden_states_output(_a      :     Dict							,_a      :     List[Any]							,_a      :     Union[str, Any]       ):
        _a				:	Optional[Any]			=   model_class(_a       )
        model.to(_a       )
        model.eval()

        with torch.no_grad():
           _a				:	Union[str, Any]			=   model(**self._prepare_for_class(_a							,_a       )       )

        _a				:	List[str]			=   outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

        _a				:	List[str]			=   self.model_tester.num_stages
        self.assertEqual(len(_a       )							,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]							,)

     _a				:	Any			=   self.model_tester.prepare_config_and_inputs_for_common()
     _a				:	int			=   ['preactivation', 'bottleneck']
     for model_class in self.all_model_classes:
        for layer_type in layers_type:
           _a				:	int			=   layer_type
           _a				:	List[str]			=   True
           check_hidden_states_output(_a							,_a							,_a       )

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           _a				:	int			=   True

           check_hidden_states_output(_a							,_a							,_a       )


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


     '''simple docstring'''






     pass


  def   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     _a				:	Optional[Any]			=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*_a       )


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


     '''simple docstring'''






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





def 					UpperCAmelCase_		():






   """simple docstring"""

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







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





  """simple docstring"""

  @cached_property
  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






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


  @slow
  def   __lowercase		(		self      :     Dict       ):


     '''simple docstring'''






     _a				:	Any			=   BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]       ).to(_a       )

     _a				:	Tuple			=   self.default_image_processor
     _a				:	str			=   prepare_img()
     _a				:	Tuple			=   image_processor(images=_a							,return_tensors='pt'       ).to(_a       )

     # forward pass
     with torch.no_grad():
        _a				:	Any			=   model(**_a       )

     # verify the logits
     _a				:	Any			=   torch.Size((1, 1000)       )
     self.assertEqual(outputs.logits.shape							,_a       )

     _a				:	Optional[Any]			=   torch.tensor([[-0.6526, -0.5263, -1.4398]]       ).to(_a       )

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







@require_torch
class UpperCAmelCase__	( __UpperCAmelCase    ,				unittest.TestCase	):





  """simple docstring"""

  __UpperCAmelCase				:				Optional[int]   =					(BitBackbone,) if is_torch_available() else ()
  __UpperCAmelCase				:				Optional[Any]   =					BitConfig

  __UpperCAmelCase				:				int   =					False


  def   __lowercase		(		self      :     int       ):


     '''simple docstring'''






     _a				:	str			=   BitModelTester(self       )










from dataclasses import dataclass, field
from typing import Tuple

from ..utils import cached_property, is_torch_available, is_torch_tpu_available, logging, requires_backends
from .benchmark_args_utils import BenchmarkArguments


if is_torch_available():
    import torch

if is_torch_tpu_available(check_device=False):
    import torch_xla.core.xla_model as xm


lowercase      :       List[str]									=  logging.get_logger(__name__)





@dataclass
class  A__    (							__UpperCAmelCase	):




   """simple docstring"""



   __A   :				Optional[Any]         =  [
       '''no_inference''',
       '''no_cuda''',
       '''no_tpu''',
       '''no_speed''',
       '''no_memory''',
       '''no_env_print''',
       '''no_multi_process''',
   ]

   def __init__(					self					,  **lowercase)			->						Tuple:







          '''simple docstring'''

          for deprecated_arg in self.deprecated_args:
                 if deprecated_arg in kwargs:
                        a__						:	Union[str, Any] =				deprecated_arg[3:]
                        setattr(self					,  lowercase					,  not kwargs.pop(lowercase))
                        logger.warning(
                            F'{deprecated_arg} is depreciated. Please use --no_{positive_arg} or'
                            F' {positive_arg}={kwargs[positive_arg]}')

          a__						:	Union[str, Any] =				kwargs.pop('torchscript'					,  self.torchscript)
          a__						:	Tuple =				kwargs.pop('torch_xla_tpu_print_metrics'					,  self.torch_xla_tpu_print_metrics)
          a__						:	Tuple =				kwargs.pop('fp16_opt_level'					,  self.fpaa_opt_level)
          super().__init__(**lowercase)

   __A   :				bool         =  field(default=__UpperCAmelCase ,  metadata={'''help''': '''Trace the models using torchscript'''}	)
   __A   :				bool         =  field(default=__UpperCAmelCase ,  metadata={'''help''': '''Print Xla/PyTorch tpu metrics'''}	)
   __A   :				str         =  field(
       default='''O1''' ,  metadata={
           '''help''': (
               '''For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\']. '''
               '''See details at https://nvidia.github.io/apex/amp.html'''
           )
       } ,  )

   @cached_property
   def 							__lowercase		(					self)			->						Tuple["torch.device", int]:







          '''simple docstring'''

          requires_backends(self					,  ['torch'])
          logger.info('PyTorch: setting up devices')
          if not self.cuda:
                 a__						:	List[str] =				torch.device('cpu')
                 a__						:	Optional[Any] =				0
          elif is_torch_tpu_available():
                 a__						:	List[str] =				xm.xla_device()
                 a__						:	Union[str, Any] =				0
          else:
                 a__						:	List[Any] =				torch.device('cuda' if torch.cuda.is_available() else 'cpu')
                 a__						:	Optional[int] =				torch.cuda.device_count()
          return device, n_gpu

   @property
   def 							__lowercase		(					self)			->						List[str]:







          '''simple docstring'''

          return is_torch_tpu_available() and self.tpu

   @property
   def 							__lowercase		(					self)			->						int:







          '''simple docstring'''

          requires_backends(self					,  ['torch'])
          # TODO(PVP): currently only single GPU is supported
          return torch.cuda.current_device()

   @property
   def 							__lowercase		(					self)			->						"torch.device":







          '''simple docstring'''

          requires_backends(self					,  ['torch'])
          return self._setup_devices[0]

   @property
   def 							__lowercase		(					self)			->						Dict:







          '''simple docstring'''

          requires_backends(self					,  ['torch'])
          return self._setup_devices[1]

   @property
   def 							__lowercase		(					self)			->						Optional[int]:







          '''simple docstring'''

          return self.n_gpu > 0


