






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




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


            """simple docstring"""







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




       def 					a_					(					self      : str					):


            """simple docstring"""







            __lowerCamelCase						:       Union[str, Any]				   =				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

            __lowerCamelCase						:       List[str]				   =				None
            if self.use_labels:
                 __lowerCamelCase						:       List[Any]				   =				ids_tensor([self.batch_size]	,    self.num_labels					)

            __lowerCamelCase						:       Dict				   =				self.get_config()

            return config, pixel_values, labels




       def 					a_					(					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 					a_					(					self      : Any	,    A__      : Dict	,    A__      : Optional[Any]	,    A__      : List[str]					):


            """simple docstring"""







            __lowerCamelCase						:       int				   =				BitModel(config=A__					)
            model.to(A__					)
            model.eval()
            __lowerCamelCase						:       Union[str, Any]				   =				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 					a_					(					self      : Optional[Any]	,    A__      : Optional[int]	,    A__      : str	,    A__      : Dict					):


            """simple docstring"""







            __lowerCamelCase						:       List[Any]				   =				self.num_labels
            __lowerCamelCase						:       int				   =				BitForImageClassification(A__					)
            model.to(A__					)
            model.eval()
            __lowerCamelCase						:       int				   =				model(A__	,    labels=A__					)
            self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.num_labels)					)




       def 					a_					(					self      : Any	,    A__      : Tuple	,    A__      : Optional[int]	,    A__      : Optional[Any]					):


            """simple docstring"""







            __lowerCamelCase						:       Any				   =				BitBackbone(config=A__					)
            model.to(A__					)
            model.eval()
            __lowerCamelCase						:       Dict				   =				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
            __lowerCamelCase						:       str				   =				None
            __lowerCamelCase						:       Union[str, Any]				   =				BitBackbone(config=A__					)
            model.to(A__					)
            model.eval()
            __lowerCamelCase						:       List[str]				   =				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 					a_					(					self      : Any					):


            """simple docstring"""







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







@require_torch
class 	SCREAMING_SNAKE_CASE	(      lowerCAmelCase_    ,   lowerCAmelCase_    ,   unittest.TestCase  ):
       snake_case__    :     List[str]   		=			(BitModel, BitForImageClassification, BitBackbone) if is_torch_available() else ()
       snake_case__    :     Union[str, Any]   		=			(
           {'feature-extraction': BitModel, 'image-classification': BitForImageClassification}
           if is_torch_available()
           else {}
       )

       snake_case__    :     Optional[int]   		=			False
       snake_case__    :     Dict   		=			False
       snake_case__    :     str   		=			False
       snake_case__    :     Optional[int]   		=			False
       snake_case__    :     Tuple   		=			False




       def 					a_					(					self      : int					):


            """simple docstring"""







            __lowerCamelCase						:       Optional[int]				   =				BitModelTester(self					)
            __lowerCamelCase						:       Optional[int]				   =				ConfigTester(self	,    config_class=A__	,    has_text_modality=A__					)




       def 					a_					(					self      : Optional[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 					a_					(					self      : Tuple					):


            """simple docstring"""







            return




       @unittest.skip(reason="""Bit does not output attentions"""					)
       def 					a_					(					self      : List[str]					):


            """simple docstring"""







            pass




       @unittest.skip(reason="""Bit does not use inputs_embeds"""					)
       def 					a_					(					self      : str					):


            """simple docstring"""







            pass




       @unittest.skip(reason="""Bit does not support input and output embeddings"""					)
       def 					a_					(					self      : Dict					):


            """simple docstring"""







            pass




       def 					a_					(					self      : Dict					):


            """simple docstring"""







            __lowerCamelCase   ,      __lowerCamelCase						:       str				   =				self.model_tester.prepare_config_and_inputs_for_common()

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

                 __lowerCamelCase						:       Tuple				   =				["""pixel_values"""]
                 self.assertListEqual(arg_names[:1]	,    A__					)




       def 					a_					(					self      : List[Any]					):


            """simple docstring"""







            __lowerCamelCase						:       Tuple				   =				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*A__					)




       def 					a_					(					self      : Dict					):


            """simple docstring"""







            __lowerCamelCase						:       Tuple				   =				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_backbone(*A__					)




       def 					a_					(					self      : int					):


            """simple docstring"""







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

            for model_class in self.all_model_classes:
                 __lowerCamelCase						:       Optional[int]				   =				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 					a_					(					self      : Any					):


            """simple docstring"""







            def check_hidden_states_output(A__      : Optional[int]	,    A__      : Dict	,    A__      : Optional[Any]					):
                 __lowerCamelCase						:       Any				   =				model_class(A__					)
                 model.to(A__					)
                 model.eval()

                 with torch.no_grad():
                      __lowerCamelCase						:       str				   =				model(**self._prepare_for_class(A__	,    A__					)					)

                 __lowerCamelCase						:       List[Any]				   =				outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                 __lowerCamelCase						:       List[Any]				   =				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]	,    )

            __lowerCamelCase   ,      __lowerCamelCase						:       List[Any]				   =				self.model_tester.prepare_config_and_inputs_for_common()
            __lowerCamelCase						:       int				   =				["""preactivation""", """bottleneck"""]
            for model_class in self.all_model_classes:
                 for layer_type in layers_type:
                      __lowerCamelCase						:       Tuple				   =				layer_type
                      __lowerCamelCase						:       Tuple				   =				True
                      check_hidden_states_output(A__	,    A__	,    A__					)

                      # check that output_hidden_states also work using config
                      del inputs_dict["output_hidden_states"]
                      __lowerCamelCase						:       List[str]				   =				True

                      check_hidden_states_output(A__	,    A__	,    A__					)




       @unittest.skip(reason="""Bit does not use feedforward chunking"""					)
       def 					a_					(					self      : int					):


            """simple docstring"""







            pass




       def 					a_					(					self      : Optional[Any]					):


            """simple docstring"""







            __lowerCamelCase						:       str				   =				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_image_classification(*A__					)




       @slow
       def 					a_					(					self      : Any					):


            """simple docstring"""







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

def     __lowercase  ()   ->     List[Any]:


     """simple docstring"""
     __lowerCamelCase						:       Tuple				   =				Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""     )
     return image







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




       @cached_property
       def 					a_					(					self      : List[Any]					):


            """simple docstring"""







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




       @slow
       def 					a_					(					self      : Dict					):


            """simple docstring"""







            __lowerCamelCase						:       Optional[int]				   =				BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]					).to(A__					)

            __lowerCamelCase						:       List[str]				   =				self.default_image_processor
            __lowerCamelCase						:       Optional[Any]				   =				prepare_img()
            __lowerCamelCase						:       List[str]				   =				image_processor(images=A__	,    return_tensors="""pt"""					).to(A__					)

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

            # verify the logits
            __lowerCamelCase						:       Tuple				   =				torch.Size((1, 1000)					)
            self.assertEqual(outputs.logits.shape	,    A__					)

            __lowerCamelCase						:       List[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 	SCREAMING_SNAKE_CASE	(      lowerCAmelCase_    ,   unittest.TestCase  ):
       snake_case__    :     Dict   		=			(BitBackbone,) if is_torch_available() else ()
       snake_case__    :     Optional[int]   		=			BitConfig

       snake_case__    :     str   		=			False




       def 					a_					(					self      : Tuple					):


            """simple docstring"""







            __lowerCamelCase						:       List[str]				   =				BitModelTester(self					)











'''simple docstring'''

import argparse

import torch

from transformers import RemBertConfig, RemBertModel, load_tf_weights_in_rembert
from transformers.utils import logging


logging.set_verbosity_info()
def     __lowercase  (_lowercase,						_lowercase,						_lowercase     )   ->     Optional[Any]:


     """simple docstring"""
     # Initialise PyTorch model
     __lowerCamelCase						:       str				   =				RemBertConfig.from_json_file(_lowercase     )
     print("""Building PyTorch model from configuration: {}""".format(str(_lowercase     )     )     )
     __lowerCamelCase						:       List[Any]				   =				RemBertModel(_lowercase     )

     # Load weights from tf checkpoint
     load_tf_weights_in_rembert(_lowercase,						_lowercase,						_lowercase     )

     # Save pytorch-model
     print("""Save PyTorch model to {}""".format(_lowercase     )     )
     torch.save(model.state_dict(),						_lowercase     )


if __name__ == "__main__":
   UpperCAmelCase__						:Tuple          =			argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       """--tf_checkpoint_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
   )
   parser.add_argument(
       """--rembert_config_file""",
       default=None,
       type=str,
       required=True,
       help=(
           """The config json file corresponding to the pre-trained RemBERT model. \n"""
           """This specifies the model architecture."""
       ),
   )
   parser.add_argument(
       """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
   )
   UpperCAmelCase__						:List[Any]          =			parser.parse_args()
   convert_rembert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.rembert_config_file, args.pytorch_dump_path)



