



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



 '''simple docstring'''







 def __init__(self					,  A					,  A=3					,  A=32					,  A=3					,  A=10					,  A=[8, 16, 32, 64]					,  A=[1, 1, 2, 1]					,  A=True					,  A=True					,  A="relu"					,  A=3					,  A=None					,  A=["stage2", "stage3", "stage4"]					,  A=[2, 3, 4]					,  A=1					,  )       ->	Any:


      """simple docstring"""



      _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(A			)
      _a           =    out_features
      _a           =    out_indices
      _a           =    num_groups



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


      """simple docstring"""



      _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    a__    (self			)       ->	Union[str, 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    a__    (self					,  A					,  A					,  A			)       ->	List[str]:


      """simple docstring"""



      _a           =    BitModel(config=A			)
      model.to(A			)
      model.eval()
      _a           =    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					,  A					,  A					,  A			)       ->	List[Any]:


      """simple docstring"""



      _a           =    self.num_labels
      _a           =    BitForImageClassification(A			)
      model.to(A			)
      model.eval()
      _a           =    model(A					,  labels=A			)
      self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.num_labels)			)



 def    a__    (self					,  A					,  A					,  A			)       ->	Any:


      """simple docstring"""



      _a           =    BitBackbone(config=A			)
      model.to(A			)
      model.eval()
      _a           =    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           =    None
      _a           =    BitBackbone(config=A			)
      model.to(A			)
      model.eval()
      _a           =    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"""



      _a           =    self.prepare_config_and_inputs()
      _a    ,       _a    ,       _a           =    config_and_inputs
      _a           =    {'''pixel_values''': pixel_values}
      return config, inputs_dict






@require_torch
class __A     (    A  ,  A  ,  unittest.TestCase		):



 '''simple docstring'''







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

 __lowerCamelCase       :			Optional[int]		=			False
 __lowerCamelCase       :			int		=			False
 __lowerCamelCase       :			Optional[int]		=			False
 __lowerCamelCase       :			Union[str, Any]		=			False
 __lowerCamelCase       :			Optional[Any]		=			False



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


      """simple docstring"""



      _a           =    BitModelTester(self			)
      _a           =    ConfigTester(self					,  config_class=A					,  has_text_modality=A			)



 def    a__    (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    a__    (self			)       ->	Union[str, Any]:


      """simple docstring"""



      return



 @unittest.skip(reason='''Bit does not output attentions'''			)
 def    a__    (self			)       ->	int:


      """simple docstring"""



      pass



 @unittest.skip(reason='''Bit does not use inputs_embeds'''			)
 def    a__    (self			)       ->	Any:


      """simple docstring"""



      pass



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


      """simple docstring"""



      pass



 def    a__    (self			)       ->	Optional[Any]:


      """simple docstring"""



      _a    ,       _a           =    self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           _a           =    model_class(A			)
           _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]					,  A			)



 def    a__    (self			)       ->	List[str]:


      """simple docstring"""



      _a           =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*A			)



 def    a__    (self			)       ->	Optional[int]:


      """simple docstring"""



      _a           =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_backbone(*A			)



 def    a__    (self			)       ->	Any:


      """simple docstring"""



      _a    ,       _a           =    self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           _a           =    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			)       ->	Tuple:


      """simple docstring"""



      def check_hidden_states_output(A					,  A					,  A			):
           _a           =    model_class(A			)
           model.to(A			)
           model.eval()

           with torch.no_grad():
                _a           =    model(**self._prepare_for_class(A					,  A			)			)

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

           _a           =    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    ,       _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(A					,  A					,  A			)

                # check that output_hidden_states also work using config
                del inputs_dict["output_hidden_states"]
                _a           =    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"""



      _a           =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*A			)



 @slow
 def    a__    (self			)       ->	int:


      """simple docstring"""



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





def lowerCAmelCase    ():



     """simple docstring"""



     _a           =    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			)       ->	str:


      """simple docstring"""



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



 @slow
 def    a__    (self			)       ->	int:


      """simple docstring"""



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

      _a           =    self.default_image_processor
      _a           =    prepare_img()
      _a           =    image_processor(images=A					,  return_tensors='''pt'''			).to(A			)

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

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

      _a           =    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 __A     (    A  ,  unittest.TestCase		):



 '''simple docstring'''







 __lowerCamelCase       :			List[Any]		=			(BitBackbone,) if is_torch_available() else ()
 __lowerCamelCase       :			List[str]		=			BitConfig

 __lowerCamelCase       :			Any		=			False



 def    a__    (self			)       ->	Dict:


      """simple docstring"""



      _a           =    BitModelTester(self			)








'''simple docstring'''





import argparse
import re

from flax.traverse_util import flatten_dict, unflatten_dict
from tax import checkpoints

from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration
from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model
from transformers.utils import logging


logging.set_verbosity_info()


# should not include what is already done by the `from_pt` argument
lowercase_									=    {
    "/attention/": "/0/SelfAttention/",
    "/self_attention/": "/0/SelfAttention/",
    "/encoder_decoder_attention/": "/1/EncDecAttention/",
    "value": "v",
    "query": "q",
    "key": "k",
    "out": "o",
    "pre_self_attention_layer_norm": "0/layer_norm",
    "pre_cross_attention_layer_norm": "1/layer_norm",
    "pre_attention_layer_norm": "0/layer_norm",  # previously 1, but seems wrong
    "token_embedder": "shared",
    "encoder_norm": "final_layer_norm",
    "decoder_norm": "final_layer_norm",
    "relpos_bias/rel_embedding": "block/0/layer/0/SelfAttention/relative_attention_bias/weight",
    "router/router_weights/w/": "router/classifier/",
    "roer/roer_weights/w/": "router/classifier/",
    "logits_dense": "lm_head",
}




def lowerCAmelCase    (__A):



     """simple docstring"""



     _a           =    list(s_dict.keys())
     for key in keys:
          _a           =    r'''.*/layers_(\d+)'''
          _a           =    key
          if re.match(__A    ,      __A):
               _a           =    re.sub(r'''layers_(\d+)'''    ,      r'''block/\1/layer'''    ,      __A)

          _a           =    r'''(encoder|decoder)\/'''

          if re.match(__A    ,      __A):
               _a           =    re.match(__A    ,      __A).groups()
               if groups[0] == "encoder":
                    _a           =    re.sub(r'''/mlp/'''    ,      r'''/1/mlp/'''    ,      __A)
                    _a           =    re.sub(r'''/pre_mlp_layer_norm/'''    ,      r'''/1/layer_norm/'''    ,      __A)

               elif groups[0] == "decoder":
                    _a           =    re.sub(r'''/mlp/'''    ,      r'''/2/mlp/'''    ,      __A)
                    _a           =    re.sub(r'''/pre_mlp_layer_norm/'''    ,      r'''/2/layer_norm/'''    ,      __A)

        # 2. Convert other classic mappings
          for old_key, temp_key in MOE_LAYER_NAME_MAPPING.items():
               if old_key in new_key:
                    _a           =    new_key.replace(__A    ,      __A)

          print(F'''{key} -> {new_key}''')
          _a           =    s_dict.pop(__A)

     if "encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
          _a           =    s_dict[
              '''encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight'''
          ].T
     if "decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
          _a           =    s_dict[
              '''decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight'''
          ].T

     # 3. Take extra care of the EXPERTS layer
     for key in list(s_dict.keys()):
          if "expert" in key:
               _a           =    s_dict[key].shape[0]
               _a           =    s_dict[key]
               for idx in range(__A):
                    _a           =    expert_weihts[idx]
                    print(F'''{key} -> {key.replace('expert/'    ,      'nested fstring')}''')

               s_dict.pop(__A)

     return s_dict


lowercase_									=    {
    "NUM_ENCODER_LAYERS": "num_layers",
    "NUM_DECODER_LAYERS": "num_decoder_layers",
    "NUM_HEADS": "num_heads",
    "HEAD_DIM": "d_kv",
    "EMBED_DIM": "d_model",
    "MLP_DIM": "d_ff",
    "NUM_SELECTED_EXPERTS": "num_selected_experts",
    "NUM_ENCODER_SPARSE_LAYERS": "num_sparse_encoder_layers",
    "NUM_DECODER_SPARSE_LAYERS": "num_sparse_decoder_layers",
    "dense.MlpBlock.activations": "feed_forward_proj",
}




def lowerCAmelCase    (__A    ,      __A):



     """simple docstring"""



     import regex as re

     with open(__A    ,      '''r''') as f:
          _a           =    f.read()

     _a           =    re.findall(r'''(.*) = ([0-9.]*)'''    ,      __A)
     _a           =    {}
     for param, value in regex_match:
          if param in GIN_TO_CONFIG_MAPPING and value != "":
               _a           =    float(__A) if '''.''' in value else int(__A)

     _a           =    re.findall(r'''(.*activations) = \(\'(.*)\',\)'''    ,      __A)[0]
     _a           =    str(activation[1])

     _a           =    num_experts
     _a           =    SwitchTransformersConfig(**__A)
     return config




def lowerCAmelCase    (__A    ,      __A    ,      __A=None    ,      __A="./"    ,      __A=8):



     """simple docstring"""



     print(F'''Loading flax weights from : {flax_checkpoint_path}''')
     _a           =    checkpoints.load_tax_checkpoint(__A)

     if gin_file is not None:
          _a           =    convert_gin_to_config(__A    ,      __A)
     else:
          _a           =    SwitchTransformersConfig.from_pretrained(__A)

     _a           =    SwitchTransformersForConditionalGeneration(__A)

     _a           =    flax_params['''target''']
     _a           =    flatten_dict(__A    ,      sep='''/''')
     _a           =    rename_keys(__A)
     _a           =    unflatten_dict(__A    ,      sep='''/''')

     # Load the flax params in the PT model
     load_flax_weights_in_pytorch_model(__A    ,      __A)

     print(F'''Save PyTorch model to {pytorch_dump_path}''')
     pt_model.save_pretrained(__A)


if __name__ == "__main__":
  lowercase_									=    argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      "--switch_t5x_checkpoint_path",
      default=None,
      type=str,
      required=True,
      help=(
          "The config json file corresponding to the pre-trained SwitchTransformers model. \nThis specifies the"
          " model architecture. If not provided, a `gin_file` has to be provided."
      ),
  )
  parser.add_argument(
      "--gin_file",
      default=None,
      type=str,
      required=False,
      help="Path to the gin config file. If not provided, a `config_file` has to be passed   ",
  )
  parser.add_argument(
      "--config_name", default=None, type=str, required=False, help="Config name of SwitchTransformers model."
  )
  parser.add_argument(
      "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output pytorch model."
  )
  parser.add_argument("--num_experts", default=8, type=int, required=False, help="Number of experts")
  lowercase_									=    parser.parse_args()
  convert_flax_checkpoint_to_pytorch(
      args.switch_tax_checkpoint_path,
      args.config_name,
      args.gin_file,
      args.pytorch_dump_folder_path,
      args.num_experts,
  )



