


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:
       def __init__(      self,   lowerCamelCase,   lowerCamelCase=3,   lowerCamelCase=32,   lowerCamelCase=3,   lowerCamelCase=10,   lowerCamelCase=[8, 16, 32, 64],   lowerCamelCase=[1, 1, 2, 1],   lowerCamelCase=True,   lowerCamelCase=True,   lowerCamelCase="relu",   lowerCamelCase=3,   lowerCamelCase=None,   lowerCamelCase=["stage2", "stage3", "stage4"],   lowerCamelCase=[2, 3, 4],   lowerCamelCase=1,   ) ->     Union[str, Any]:




           """simple docstring"""

           _lowercase :							Union[str, Any]   	=	parent
           _lowercase :							str   	=	batch_size
           _lowercase :							List[str]   	=	image_size
           _lowercase :							List[str]   	=	num_channels
           _lowercase :							List[str]   	=	embeddings_size
           _lowercase :							Optional[int]   	=	hidden_sizes
           _lowercase :							List[str]   	=	depths
           _lowercase :							Union[str, Any]   	=	is_training
           _lowercase :							str   	=	use_labels
           _lowercase :							Union[str, Any]   	=	hidden_act
           _lowercase :							List[str]   	=	num_labels
           _lowercase :							Optional[Any]   	=	scope
           _lowercase :							List[Any]   	=	len(lowerCamelCase)
           _lowercase :							Optional[int]   	=	out_features
           _lowercase :							Union[str, Any]   	=	out_indices
           _lowercase :							Optional[int]   	=	num_groups
       def     UpperCamelCase  (      self) ->     Tuple:




           """simple docstring"""

           _lowercase :							Optional[int]   	=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

           _lowercase :							Union[str, Any]   	=	None
           if self.use_labels:
               _lowercase :							int   	=	ids_tensor([self.batch_size],   self.num_labels)

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

           return config, pixel_values, labels
       def     UpperCamelCase  (      self) ->     int:




           """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,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase) ->     Optional[Any]:




           """simple docstring"""

           _lowercase :							Optional[int]   	=	BitModel(config=lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase :							int   	=	model(lowerCamelCase)
           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,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase) ->     Dict:




           """simple docstring"""

           _lowercase :							Any   	=	self.num_labels
           _lowercase :							Tuple   	=	BitForImageClassification(lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase :							Tuple   	=	model(lowerCamelCase,   labels=lowerCamelCase)
           self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.num_labels))
       def     UpperCamelCase  (      self,   lowerCamelCase,   lowerCamelCase,   lowerCamelCase) ->     Union[str, Any]:




           """simple docstring"""

           _lowercase :							Tuple   	=	BitBackbone(config=lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase :							Optional[Any]   	=	model(lowerCamelCase)

           # 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 :							Dict   	=	None
           _lowercase :							Union[str, Any]   	=	BitBackbone(config=lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase :							Union[str, Any]   	=	model(lowerCamelCase)

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




           """simple docstring"""

           _lowercase :							Union[str, Any]   	=	self.prepare_config_and_inputs()
           _lowercase	,       _lowercase	,       _lowercase :							Dict   	=	config_and_inputs
           _lowercase :							List[Any]   	=	{'pixel_values': pixel_values}
           return config, inputs_dict






@require_torch
class _lowerCamelCase(		_a,				_a,				unittest.TestCase     ):
       lowercase_				:			Any									=    (BitModel, BitForImageClassification, BitBackbone) if is_torch_available() else ()
       lowercase_				:			Tuple									=    (
           {"feature-extraction": BitModel, "image-classification": BitForImageClassification}
           if is_torch_available()
           else {}
       )

       lowercase_				:			Union[str, Any]									=    False
       lowercase_				:			Any									=    False
       lowercase_				:			Any									=    False
       lowercase_				:			int									=    False
       lowercase_				:			List[str]									=    False
       def     UpperCamelCase  (      self) ->     Tuple:




           """simple docstring"""

           _lowercase :							int   	=	BitModelTester(self)
           _lowercase :							Optional[Any]   	=	ConfigTester(self,   config_class=lowerCamelCase,   has_text_modality=lowerCamelCase)
       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) ->     Optional[int]:




           """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) ->     Optional[Any]:




           """simple docstring"""

           pass
       @unittest.skip(reason='Bit does not support input and output embeddings')
       def     UpperCamelCase  (      self) ->     Dict:




           """simple docstring"""

           pass
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           _lowercase	,       _lowercase :							int   	=	self.model_tester.prepare_config_and_inputs_for_common()

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

               _lowercase :							Dict   	=	['pixel_values']
               self.assertListEqual(arg_names[:1],   lowerCamelCase)
       def     UpperCamelCase  (      self) ->     Union[str, Any]:




           """simple docstring"""

           _lowercase :							int   	=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*lowerCamelCase)
       def     UpperCamelCase  (      self) ->     List[str]:




           """simple docstring"""

           _lowercase :							Union[str, Any]   	=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_backbone(*lowerCamelCase)
       def     UpperCamelCase  (      self) ->     Any:




           """simple docstring"""

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

           for model_class in self.all_model_classes:
               _lowercase :							Dict   	=	model_class(config=lowerCamelCase)
               for name, module in model.named_modules():
                   if isinstance(lowerCamelCase,   (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) ->     str:




           """simple docstring"""

           def check_hidden_states_output(lowerCamelCase,   lowerCamelCase,   lowerCamelCase):
               _lowercase :							str   	=	model_class(lowerCamelCase)
               model.to(lowerCamelCase)
               model.eval()

               with torch.no_grad():
                   _lowercase :							Tuple   	=	model(**self._prepare_for_class(lowerCamelCase,   lowerCamelCase))

               _lowercase :							str   	=	outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

               _lowercase :							Union[str, Any]   	=	self.model_tester.num_stages
               self.assertEqual(len(lowerCamelCase),   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 :							Any   	=	self.model_tester.prepare_config_and_inputs_for_common()
           _lowercase :							Optional[Any]   	=	['preactivation', 'bottleneck']
           for model_class in self.all_model_classes:
               for layer_type in layers_type:
                   _lowercase :							int   	=	layer_type
                   _lowercase :							int   	=	True
                   check_hidden_states_output(lowerCamelCase,   lowerCamelCase,   lowerCamelCase)

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

                   check_hidden_states_output(lowerCamelCase,   lowerCamelCase,   lowerCamelCase)
       @unittest.skip(reason='Bit does not use feedforward chunking')
       def     UpperCamelCase  (      self) ->     Tuple:




           """simple docstring"""

           pass
       def     UpperCamelCase  (      self) ->     Union[str, Any]:




           """simple docstring"""

           _lowercase :							str   	=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*lowerCamelCase)




       @slow
       def     UpperCamelCase  (      self) ->     str:




           """simple docstring"""

           for model_name in BIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _lowercase :							str   	=	BitModel.from_pretrained(lowerCamelCase)
               self.assertIsNotNone(lowerCamelCase)



def        UpperCamelCase_(   )					->		List[Any]:
    _lowercase :							Optional[int]   	=	Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'				)
    return image






@require_torch
@require_vision
class _lowerCamelCase(		unittest.TestCase     ):
       @cached_property
       def     UpperCamelCase  (      self) ->     List[Any]:




           """simple docstring"""

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




       @slow
       def     UpperCamelCase  (      self) ->     str:




           """simple docstring"""

           _lowercase :							Optional[Any]   	=	BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]).to(lowerCamelCase)

           _lowercase :							List[str]   	=	self.default_image_processor
           _lowercase :							int   	=	prepare_img()
           _lowercase :							Any   	=	image_processor(images=lowerCamelCase,   return_tensors='pt').to(lowerCamelCase)

           # forward pass
           with torch.no_grad():
               _lowercase :							int   	=	model(**lowerCamelCase)

           # verify the logits
           _lowercase :							Dict   	=	torch.Size((1, 10_00))
           self.assertEqual(outputs.logits.shape,   lowerCamelCase)

           _lowercase :							Optional[Any]   	=	torch.tensor([[-0.6_5_2_6, -0.5_2_6_3, -1.4_3_9_8]]).to(lowerCamelCase)

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






@require_torch
class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			Optional[Any]									=    (BitBackbone,) if is_torch_available() else ()
       lowercase_				:			int									=    BitConfig

       lowercase_				:			Union[str, Any]									=    False
       def     UpperCamelCase  (      self) ->     Any:




           """simple docstring"""

           _lowercase :							Tuple   	=	BitModelTester(self)






import dataclasses
import json
import warnings
from dataclasses import dataclass, field
from time import time
from typing import List

from ..utils import logging


SCREAMING_SNAKE_CASE :   Dict							      =				logging.get_logger(__name__)



def        UpperCamelCase_(   lowerCamelCase_=None							,				lowerCamelCase_=None				)					->		Tuple:
    return field(default_factory=lambda: default							,				metadata=lowerCamelCase_				)






@dataclass
class _lowerCamelCase:
       lowercase_				:			List[str]									=    list_field(
           default=[],				metadata={
               """help""": (
                   """Model checkpoints to be provided to the AutoModel classes. Leave blank to benchmark the base version"""
                   """ of all available models"""
               )
           },				)

       lowercase_				:			List[int]									=    list_field(
           default=[8],				metadata={"""help""": """List of batch sizes for which memory and time performance will be evaluated"""}     )

       lowercase_				:			List[int]									=    list_field(
           default=[8, 32, 1_28, 5_12],				metadata={"""help""": """List of sequence lengths for which memory and time performance will be evaluated"""},				)

       lowercase_				:			bool									=    field(
           default=_a,				metadata={"""help""": """Whether to benchmark inference of model. Inference can be disabled via --no-inference."""},				)
       lowercase_				:			bool									=    field(
           default=_a,				metadata={"""help""": """Whether to run on available cuda devices. Cuda can be disabled via --no-cuda."""},				)
       lowercase_				:			bool									=    field(
           default=_a,				metadata={"""help""": """Whether to run on available tpu devices. TPU can be disabled via --no-tpu."""}     )
       lowercase_				:			bool									=    field(default=_a,				metadata={"""help""": """Use FP16 to accelerate inference."""}     )
       lowercase_				:			bool									=    field(default=_a,				metadata={"""help""": """Benchmark training of model"""}     )
       lowercase_				:			bool									=    field(default=_a,				metadata={"""help""": """Verbose memory tracing"""}     )
       lowercase_				:			bool									=    field(
           default=_a,				metadata={"""help""": """Whether to perform speed measurements. Speed measurements can be disabled via --no-speed."""},				)
       lowercase_				:			bool									=    field(
           default=_a,				metadata={
               """help""": """Whether to perform memory measurements. Memory measurements can be disabled via --no-memory"""
           },				)
       lowercase_				:			bool									=    field(default=_a,				metadata={"""help""": """Trace memory line by line"""}     )
       lowercase_				:			bool									=    field(default=_a,				metadata={"""help""": """Save result to a CSV file"""}     )
       lowercase_				:			bool									=    field(default=_a,				metadata={"""help""": """Save all print statements in a log file"""}     )
       lowercase_				:			bool									=    field(default=_a,				metadata={"""help""": """Whether to print environment information"""}     )
       lowercase_				:			bool									=    field(
           default=_a,				metadata={
               """help""": (
                   """Whether to use multiprocessing for memory and speed measurement. It is highly recommended to use"""
                   """ multiprocessing for accurate CPU and GPU memory measurements. This option should only be disabled"""
                   """ for debugging / testing and on TPU."""
               )
           },				)
       lowercase_				:			str									=    field(
           default=F'''inference_time_{round(time()     )}.csv''',				metadata={"""help""": """CSV filename used if saving time results to csv."""},				)
       lowercase_				:			str									=    field(
           default=F'''inference_memory_{round(time()     )}.csv''',				metadata={"""help""": """CSV filename used if saving memory results to csv."""},				)
       lowercase_				:			str									=    field(
           default=F'''train_time_{round(time()     )}.csv''',				metadata={"""help""": """CSV filename used if saving time results to csv for training."""},				)
       lowercase_				:			str									=    field(
           default=F'''train_memory_{round(time()     )}.csv''',				metadata={"""help""": """CSV filename used if saving memory results to csv for training."""},				)
       lowercase_				:			str									=    field(
           default=F'''env_info_{round(time()     )}.csv''',				metadata={"""help""": """CSV filename used if saving environment information."""},				)
       lowercase_				:			str									=    field(
           default=F'''log_{round(time()     )}.csv''',				metadata={"""help""": """Log filename used if print statements are saved in log."""},				)
       lowercase_				:			int									=    field(default=3,				metadata={"""help""": """Times an experiment will be run."""}     )
       lowercase_				:			bool									=    field(
           default=_a,				metadata={
               """help""": (
                   """Instead of loading the model as defined in `config.architectures` if exists, just load the pretrain"""
                   """ model weights."""
               )
           },				)
       def     UpperCamelCase  (      self) ->     Union[str, Any]:




           """simple docstring"""

           warnings.warn(
               F'''The class {self.__class__} is deprecated. Hugging Face Benchmarking utils'''
               ' are deprecated in general and it is advised to use external Benchmarking libraries '
               ' to benchmark Transformer models.',   lowerCamelCase,   )
       def     UpperCamelCase  (      self) ->     Optional[Any]:




           """simple docstring"""

           return json.dumps(dataclasses.asdict(self),   indent=2)
       @property
       def     UpperCamelCase  (      self) ->     List[str]:




           """simple docstring"""

           if len(self.models) <= 0:
               raise ValueError(
                   'Please make sure you provide at least one model name / model identifier, *e.g.* `--models'
                   ' bert-base-cased` or `args.models = [\'bert-base-cased\'].')
           return self.models




       @property
       def     UpperCamelCase  (      self) ->     Dict:




           """simple docstring"""

           if not self.multi_process:
               return False
           elif self.is_tpu:
               logger.info('Multiprocessing is currently not possible on TPU.')
               return False
           else:
               return True


