


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     A__   :







    """simple docstring"""

    def __init__(			self     ,   __snake_case     ,   __snake_case=3     ,   __snake_case=3_2     ,   __snake_case=3     ,   __snake_case=1_0     ,   __snake_case=[8, 1_6, 3_2, 6_4]     ,   __snake_case=[1, 1, 2, 1]     ,   __snake_case=True     ,   __snake_case=True     ,   __snake_case="relu"     ,   __snake_case=3     ,   __snake_case=None     ,   __snake_case=["stage2", "stage3", "stage4"]     ,   __snake_case=[2, 3, 4]     ,   __snake_case=1     ,   ):
         snake_case											=							parent
         snake_case											=							batch_size
         snake_case											=							image_size
         snake_case											=							num_channels
         snake_case											=							embeddings_size
         snake_case											=							hidden_sizes
         snake_case											=							depths
         snake_case											=							is_training
         snake_case											=							use_labels
         snake_case											=							hidden_act
         snake_case											=							num_labels
         snake_case											=							scope
         snake_case											=							len(__snake_case       )
         snake_case											=							out_features
         snake_case											=							out_indices
         snake_case											=							num_groups
    def 						a_     (			self       ):
         snake_case											=							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

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

         snake_case											=							self.get_config()

         return config, pixel_values, labels
    def 						a_     (			self       ):
         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 						a_     (			self     ,   __snake_case     ,   __snake_case     ,   __snake_case       ):
         snake_case											=							BitModel(config=__snake_case       )
         model.to(__snake_case       )
         model.eval()
         snake_case											=							model(__snake_case       )
         self.parent.assertEqual(
             result.last_hidden_state.shape     ,   (self.batch_size, self.hidden_sizes[-1], self.image_size // 3_2, self.image_size // 3_2)     ,   )
    def 						a_     (			self     ,   __snake_case     ,   __snake_case     ,   __snake_case       ):
         snake_case											=							self.num_labels
         snake_case											=							BitForImageClassification(__snake_case       )
         model.to(__snake_case       )
         model.eval()
         snake_case											=							model(__snake_case     ,   labels=__snake_case       )
         self.parent.assertEqual(result.logits.shape     ,   (self.batch_size, self.num_labels)       )
    def 						a_     (			self     ,   __snake_case     ,   __snake_case     ,   __snake_case       ):
         snake_case											=							BitBackbone(config=__snake_case       )
         model.to(__snake_case       )
         model.eval()
         snake_case											=							model(__snake_case       )

         # 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
         snake_case											=							None
         snake_case											=							BitBackbone(config=__snake_case       )
         model.to(__snake_case       )
         model.eval()
         snake_case											=							model(__snake_case       )

         # 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 						a_     (			self       ):
         snake_case											=							self.prepare_config_and_inputs()
         snake_case					,    snake_case					,    snake_case											=							config_and_inputs
         snake_case											=							{'''pixel_values''': pixel_values}
         return config, inputs_dict






@require_torch
class     A__   (       snake_case__			,   snake_case__			,   unittest.TestCase      ):







    """simple docstring"""

    __magic_name__         =						(BitModel, BitForImageClassification, BitBackbone) if is_torch_available() else ()
    __magic_name__         =						(
        {'feature-extraction': BitModel, 'image-classification': BitForImageClassification}
        if is_torch_available()
        else {}
    )

    __magic_name__         =						False
    __magic_name__         =						False
    __magic_name__         =						False
    __magic_name__         =						False
    __magic_name__         =						False
    def 						a_     (			self       ):
         snake_case											=							BitModelTester(self       )
         snake_case											=							ConfigTester(self     ,   config_class=__snake_case     ,   has_text_modality=__snake_case       )
    def 						a_     (			self       ):
         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 						a_     (			self       ):
         return
    @unittest.skip(reason='''Bit does not output attentions'''       )
    def 						a_     (			self       ):
         pass
    @unittest.skip(reason='''Bit does not use inputs_embeds'''       )
    def 						a_     (			self       ):
         pass
    @unittest.skip(reason='''Bit does not support input and output embeddings'''       )
    def 						a_     (			self       ):
         pass
    def 						a_     (			self       ):
         snake_case					,    snake_case											=							self.model_tester.prepare_config_and_inputs_for_common()

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

              snake_case											=							['''pixel_values''']
              self.assertListEqual(arg_names[:1]     ,   __snake_case       )
    def 						a_     (			self       ):
         snake_case											=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__snake_case       )
    def 						a_     (			self       ):
         snake_case											=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_backbone(*__snake_case       )
    def 						a_     (			self       ):
         snake_case					,    snake_case											=							self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
              snake_case											=							model_class(config=__snake_case       )
              for name, module in model.named_modules():
                   if isinstance(__snake_case     ,   (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 						a_     (			self       ):
         def check_hidden_states_output(__snake_case     ,   __snake_case     ,   __snake_case       ):
              snake_case											=							model_class(__snake_case       )
              model.to(__snake_case       )
              model.eval()

              with torch.no_grad():
                   snake_case											=							model(**self._prepare_for_class(__snake_case     ,   __snake_case       )       )

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

              snake_case											=							self.model_tester.num_stages
              self.assertEqual(len(__snake_case       )     ,   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]     ,   )

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

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

                   check_hidden_states_output(__snake_case     ,   __snake_case     ,   __snake_case       )
    @unittest.skip(reason='''Bit does not use feedforward chunking'''       )
    def 						a_     (			self       ):
         pass
    def 						a_     (			self       ):
         snake_case											=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*__snake_case       )




    @slow
    def 						a_     (			self       ):
         for model_name in BIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              snake_case											=							BitModel.from_pretrained(__snake_case       )
              self.assertIsNotNone(__snake_case       )






def      UpperCAmelCase__		():



     """simple docstring"""






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

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







    """simple docstring"""

    @cached_property
    def 						a_     (			self       ):
         return (
             BitImageProcessor.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]       ) if is_vision_available() else None
         )




    @slow
    def 						a_     (			self       ):
         snake_case											=							BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]       ).to(__snake_case       )

         snake_case											=							self.default_image_processor
         snake_case											=							prepare_img()
         snake_case											=							image_processor(images=__snake_case     ,   return_tensors='''pt'''       ).to(__snake_case       )

         # forward pass
         with torch.no_grad():
              snake_case											=							model(**__snake_case       )

         # verify the logits
         snake_case											=							torch.Size((1, 1_0_0_0)       )
         self.assertEqual(outputs.logits.shape     ,   __snake_case       )

         snake_case											=							torch.tensor([[-0.6526, -0.5263, -1.4398]]       ).to(__snake_case       )

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






@require_torch
class     A__   (       snake_case__			,   unittest.TestCase      ):







    """simple docstring"""

    __magic_name__         =						(BitBackbone,) if is_torch_available() else ()
    __magic_name__         =						BitConfig

    __magic_name__         =						False
    def 						a_     (			self       ):
         snake_case											=							BitModelTester(self       )








from ..utils import DummyObject, requires_backends

class     A__   (       metaclass=snake_case__      ):







    """simple docstring"""

    __magic_name__         =						['flax', 'transformers']
    def __init__(			self     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(self     ,   ['''flax''', '''transformers''']       )
    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )




    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )






class     A__   (       metaclass=snake_case__      ):







    """simple docstring"""

    __magic_name__         =						['flax', 'transformers']
    def __init__(			self     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(self     ,   ['''flax''', '''transformers''']       )
    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )




    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )






class     A__   (       metaclass=snake_case__      ):







    """simple docstring"""

    __magic_name__         =						['flax', 'transformers']
    def __init__(			self     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(self     ,   ['''flax''', '''transformers''']       )
    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )




    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )






class     A__   (       metaclass=snake_case__      ):







    """simple docstring"""

    __magic_name__         =						['flax', 'transformers']
    def __init__(			self     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(self     ,   ['''flax''', '''transformers''']       )
    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )




    @classmethod
    def 						a_     (			cls     ,   *__snake_case     ,   **__snake_case       ):
         requires_backends(cls     ,   ['''flax''', '''transformers''']       )




