




"""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 							lowercase_			:



 '''simple docstring'''

 def __init__(					self			:			Tuple       ,				_UpperCAmelCase			:			int       ,				_UpperCAmelCase			:			List[Any]=3       ,				_UpperCAmelCase			:			List[Any]=32       ,				_UpperCAmelCase			:			Optional[Any]=3       ,				_UpperCAmelCase			:			int=10       ,				_UpperCAmelCase			:			int=[8, 16, 32, 64]       ,				_UpperCAmelCase			:			Optional[Any]=[1, 1, 2, 1]       ,				_UpperCAmelCase			:			Optional[Any]=True       ,				_UpperCAmelCase			:			Union[str, Any]=True       ,				_UpperCAmelCase			:			Union[str, Any]="relu"       ,				_UpperCAmelCase			:			Optional[int]=3       ,				_UpperCAmelCase			:			Tuple=None       ,				_UpperCAmelCase			:			Dict=["stage2", "stage3", "stage4"]       ,				_UpperCAmelCase			:			Union[str, Any]=[2, 3, 4]       ,				_UpperCAmelCase			:			Optional[int]=1       ,				):
       _A     =  parent
       _A     =  batch_size
       _A     =  image_size
       _A     =  num_channels
       _A     =  embeddings_size
       _A     =  hidden_sizes
       _A     =  depths
       _A     =  is_training
       _A     =  use_labels
       _A     =  hidden_act
       _A     =  num_labels
       _A     =  scope
       _A     =  len(_UpperCAmelCase						)
       _A     =  out_features
       _A     =  out_indices
       _A     =  num_groups

 def      lowerCAmelCase_				(					self			:			str						):
       _A     =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

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

       _A     =  self.get_config()

       return config, pixel_values, labels

 def      lowerCAmelCase_				(					self			:			int						):
       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      lowerCAmelCase_				(					self			:			List[Any]       ,				_UpperCAmelCase			:			Optional[Any]       ,				_UpperCAmelCase			:			Tuple       ,				_UpperCAmelCase			:			List[str]						):
       _A     =  BitModel(config=_UpperCAmelCase						)
       model.to(_UpperCAmelCase						)
       model.eval()
       _A     =  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      lowerCAmelCase_				(					self			:			Tuple       ,				_UpperCAmelCase			:			Optional[Any]       ,				_UpperCAmelCase			:			int       ,				_UpperCAmelCase			:			Tuple						):
       _A     =  self.num_labels
       _A     =  BitForImageClassification(_UpperCAmelCase						)
       model.to(_UpperCAmelCase						)
       model.eval()
       _A     =  model(_UpperCAmelCase       ,				labels=_UpperCAmelCase						)
       self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, self.num_labels)						)

 def      lowerCAmelCase_				(					self			:			List[str]       ,				_UpperCAmelCase			:			Union[str, Any]       ,				_UpperCAmelCase			:			Optional[int]       ,				_UpperCAmelCase			:			int						):
       _A     =  BitBackbone(config=_UpperCAmelCase						)
       model.to(_UpperCAmelCase						)
       model.eval()
       _A     =  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
       _A     =  None
       _A     =  BitBackbone(config=_UpperCAmelCase						)
       model.to(_UpperCAmelCase						)
       model.eval()
       _A     =  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      lowerCAmelCase_				(					self			:			Union[str, Any]						):
       _A     =  self.prepare_config_and_inputs()
       _A  ,       _A  ,       _A     =  config_and_inputs
       _A     =  {'pixel_values': pixel_values}
       return config, inputs_dict





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



 '''simple docstring'''
 UpperCAmelCase       :	str     =      (BitModel, BitForImageClassification, BitBackbone) if is_torch_available() else ()
 UpperCAmelCase       :	Tuple     =      (
     {'''feature-extraction''': BitModel, '''image-classification''': BitForImageClassification}
     if is_torch_available()
     else {}
 )

 UpperCAmelCase       :	Optional[int]     =      False
 UpperCAmelCase       :	Any     =      False
 UpperCAmelCase       :	List[str]     =      False
 UpperCAmelCase       :	List[Any]     =      False
 UpperCAmelCase       :	str     =      False

 def      lowerCAmelCase_				(					self			:			int						):
       _A     =  BitModelTester(self						)
       _A     =  ConfigTester(self       ,				config_class=_UpperCAmelCase       ,				has_text_modality=_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			Optional[Any]						):
       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      lowerCAmelCase_				(					self			:			Any						):
       return

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

 @unittest.skip(reason='Bit does not use inputs_embeds'						)
 def      lowerCAmelCase_				(					self			:			List[Any]						):
       pass

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

 def      lowerCAmelCase_				(					self			:			Optional[int]						):
       _A  ,       _A     =  self.model_tester.prepare_config_and_inputs_for_common()

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

             _A     =  ['pixel_values']
             self.assertListEqual(arg_names[:1]       ,				_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			Optional[int]						):
       _A     =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			List[Any]						):
       _A     =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_backbone(*_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			str						):
       _A  ,       _A     =  self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             _A     =  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      lowerCAmelCase_				(					self			:			List[str]						):
       def check_hidden_states_output(_UpperCAmelCase			:			Any       ,				_UpperCAmelCase			:			Tuple       ,				_UpperCAmelCase			:			Any						):
             _A     =  model_class(_UpperCAmelCase						)
             model.to(_UpperCAmelCase						)
             model.eval()

             with torch.no_grad():
                   _A     =  model(**self._prepare_for_class(_UpperCAmelCase       ,				_UpperCAmelCase						)						)

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

             _A     =  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]       ,				)

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

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

                   check_hidden_states_output(_UpperCAmelCase       ,				_UpperCAmelCase       ,				_UpperCAmelCase						)

 @unittest.skip(reason='Bit does not use feedforward chunking'						)
 def      lowerCAmelCase_				(					self			:			int						):
       pass

 def      lowerCAmelCase_				(					self			:			List[str]						):
       _A     =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*_UpperCAmelCase						)



 @slow
 def      lowerCAmelCase_				(					self			:			List[str]						):
       for model_name in BIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _A     =  BitModel.from_pretrained(_UpperCAmelCase						)
             self.assertIsNotNone(_UpperCAmelCase						)






def 		_snake_case	(   )       ->   str:






      '''simple docstring'''


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





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



 '''simple docstring'''

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



 @slow
 def      lowerCAmelCase_				(					self			:			Tuple						):
       _A     =  BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]						).to(_UpperCAmelCase						)

       _A     =  self.default_image_processor
       _A     =  prepare_img()
       _A     =  image_processor(images=_UpperCAmelCase       ,				return_tensors='pt'						).to(_UpperCAmelCase						)

       # forward pass
       with torch.no_grad():
             _A     =  model(**_UpperCAmelCase						)

       # verify the logits
       _A     =  torch.Size((1, 1_000)						)
       self.assertEqual(outputs.logits.shape       ,				_UpperCAmelCase						)

       _A     =  torch.tensor([[-0.6526, -0.5263, -1.4398]]						).to(_UpperCAmelCase						)

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





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



 '''simple docstring'''
 UpperCAmelCase       :	Tuple     =      (BitBackbone,) if is_torch_available() else ()
 UpperCAmelCase       :	List[Any]     =      BitConfig

 UpperCAmelCase       :	Optional[int]     =      False

 def      lowerCAmelCase_				(					self			:			List[Any]						):
       _A     =  BitModelTester(self						)







"""simple docstring"""

import argparse
from collections import OrderedDict
from pathlib import Path

import torch

from transformers import (
    VisualBertConfig,
    VisualBertForMultipleChoice,
    VisualBertForPreTraining,
    VisualBertForQuestionAnswering,
    VisualBertForVisualReasoning,
)
from transformers.utils import logging


logging.set_verbosity_info()
a							      =     logging.get_logger(__name__)

a							      =     [
    ('''bert.bert''', '''visual_bert'''),
    ('''bert.cls''', '''cls'''),
    ('''bert.classifier''', '''cls'''),
    ('''token_type_embeddings_visual''', '''visual_token_type_embeddings'''),
    ('''position_embeddings_visual''', '''visual_position_embeddings'''),
    ('''projection''', '''visual_projection'''),
]

a							      =     [
    '''nlvr2_coco_pre_trained.th''',
    '''nlvr2_fine_tuned.th''',
    '''nlvr2_pre_trained.th''',
    '''vcr_coco_pre_train.th''',
    '''vcr_fine_tune.th''',
    '''vcr_pre_train.th''',
    '''vqa_coco_pre_trained.th''',
    '''vqa_fine_tuned.th''',
    '''vqa_pre_trained.th''',
]






def 		_snake_case	(   _snake_case				:   Optional[Any]				)       ->   str:






      '''simple docstring'''


      _A     =  torch.load(_snake_case							,    map_location='cpu'				)
      return sd






def 		_snake_case	(   _snake_case				:   Union[str, Any]							,    _snake_case				:   str							,    _snake_case				:   Tuple=rename_keys_prefix				)       ->   List[str]:






      '''simple docstring'''


      _A     =  OrderedDict()
      _A     =  torch.arange(config.max_position_embeddings				).expand((1, -1)				)
      # detector_d = OrderedDict()
      for key in d:
            if "detector" in key:
                  # detector_d[key.replace('detector.','')] = d[key]
                  continue
            _A     =  key
            for name_pair in rename_keys_prefix:
                  _A     =  new_key.replace(name_pair[0]							,    name_pair[1]				)
            _A     =  d[key]
            if key == "bert.cls.predictions.decoder.weight":
                  # Old bert code didn't have `decoder.bias`, but was added separately
                  _A     =  new_d['cls.predictions.bias']
      return new_d






@torch.no_grad()
def 		_snake_case	(   _snake_case				:   List[str]							,    _snake_case				:   Dict				)       ->   Dict:






      '''simple docstring'''


      assert (
          checkpoint_path.split('/'				)[-1] in ACCEPTABLE_CHECKPOINTS
      ), F'''The checkpoint provided must be in {ACCEPTABLE_CHECKPOINTS}.'''

      # Get Config
      if "pre" in checkpoint_path:
            _A     =  'pretraining'
            if "vcr" in checkpoint_path:
                  _A     =  {'visual_embedding_dim': 5_12}
            elif "vqa_advanced" in checkpoint_path:
                  _A     =  {'visual_embedding_dim': 20_48}
            elif "vqa" in checkpoint_path:
                  _A     =  {'visual_embedding_dim': 20_48}
            elif "nlvr" in checkpoint_path:
                  _A     =  {'visual_embedding_dim': 10_24}
            else:
                  raise NotImplementedError(F'''No implementation found for `{checkpoint_path}`.'''				)
      else:
            if "vcr" in checkpoint_path:
                  _A     =  {'visual_embedding_dim': 5_12}
                  _A     =  'multichoice'
            elif "vqa_advanced" in checkpoint_path:
                  _A     =  {'visual_embedding_dim': 20_48}
                  _A     =  'vqa_advanced'
            elif "vqa" in checkpoint_path:
                  _A     =  {'visual_embedding_dim': 20_48, 'num_labels': 31_29}
                  _A     =  'vqa'
            elif "nlvr" in checkpoint_path:
                  _A     =  {
                      'visual_embedding_dim': 10_24,
                      'num_labels': 2,
                  }
                  _A     =  'nlvr'

      _A     =  VisualBertConfig(**_snake_case				)

      # Load State Dict
      _A     =  load_state_dict(_snake_case				)

      _A     =  get_new_dict(_snake_case							,    _snake_case				)

      if model_type == "pretraining":
            _A     =  VisualBertForPreTraining(_snake_case				)
      elif model_type == "vqa":
            _A     =  VisualBertForQuestionAnswering(_snake_case				)
      elif model_type == "nlvr":
            _A     =  VisualBertForVisualReasoning(_snake_case				)
      elif model_type == "multichoice":
            _A     =  VisualBertForMultipleChoice(_snake_case				)

      model.load_state_dict(_snake_case				)
      # Save Checkpoints
      Path(_snake_case				).mkdir(exist_ok=_snake_case				)
      model.save_pretrained(_snake_case				)


if __name__ == "__main__":
       a							      =     argparse.ArgumentParser()
       # Required parameters
       parser.add_argument('''orig_checkpoint_path''', type=str, help='''A path to .th on local filesystem.''')
       parser.add_argument('''pytorch_dump_folder_path''', type=str, help='''Path to the output PyTorch model.''')
       a							      =     parser.parse_args()
       convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)

