





'''simple docstring'''

from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import ResNetConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import TFResNetForImageClassification, TFResNetModel
    from transformers.models.resnet.modeling_tf_resnet import TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor







class __a			:





   def __init__(       self					:      Optional[int]					,lowerCamelCase					:      Dict					,lowerCamelCase					:      Any=3					,lowerCamelCase					:      Any=32					,lowerCamelCase					:      Optional[Any]=3					,lowerCamelCase					:      Any=10					,lowerCamelCase					:      Dict=[10, 20, 30, 40]					,lowerCamelCase					:      Dict=[1, 1, 2, 1]					,lowerCamelCase					:      Dict=True					,lowerCamelCase					:      Any=True					,lowerCamelCase					:      str="relu"					,lowerCamelCase					:      Optional[int]=3					,lowerCamelCase					:      List[Any]=None					,):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							parent
      __SCREAMING_SNAKE_CASE   				=							batch_size
      __SCREAMING_SNAKE_CASE   				=							image_size
      __SCREAMING_SNAKE_CASE   				=							num_channels
      __SCREAMING_SNAKE_CASE   				=							embeddings_size
      __SCREAMING_SNAKE_CASE   				=							hidden_sizes
      __SCREAMING_SNAKE_CASE   				=							depths
      __SCREAMING_SNAKE_CASE   				=							is_training
      __SCREAMING_SNAKE_CASE   				=							use_labels
      __SCREAMING_SNAKE_CASE   				=							hidden_act
      __SCREAMING_SNAKE_CASE   				=							num_labels
      __SCREAMING_SNAKE_CASE   				=							scope
      __SCREAMING_SNAKE_CASE   				=							len(lowerCamelCase     )





   def   UpperCAmelCase__  (       self					:      Optional[Any]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

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

      __SCREAMING_SNAKE_CASE   				=							self.get_config()

      return config, pixel_values, labels





   def   UpperCAmelCase__  (       self					:      Optional[Any]     ):

      '''simple docstring'''

      return ResNetConfig(
          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					,image_size=self.image_size					,)





   def   UpperCAmelCase__  (       self					:      Dict					,lowerCamelCase					:      str					,lowerCamelCase					:      List[Any]					,lowerCamelCase					:      int     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							TFResNetModel(config=lowerCamelCase     )
      __SCREAMING_SNAKE_CASE   				=							model(lowerCamelCase     )
      # expected last hidden states: B, C, H // 32, W // 32
      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					:      int					,lowerCamelCase					:      List[str]					,lowerCamelCase					:      Optional[int]					,lowerCamelCase					:      List[Any]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							self.num_labels
      __SCREAMING_SNAKE_CASE   				=							TFResNetForImageClassification(lowerCamelCase     )
      __SCREAMING_SNAKE_CASE   				=							model(lowerCamelCase					,labels=lowerCamelCase     )
      self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.num_labels)     )





   def   UpperCAmelCase__  (       self					:      List[Any]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							self.prepare_config_and_inputs()
      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE   				=							config_and_inputs
      __SCREAMING_SNAKE_CASE   				=							{"""pixel_values""": pixel_values}
      return config, inputs_dict







@require_tf
class __a			( _snake_case, _snake_case, unittest.TestCase      ):
   __UpperCamelCase :			int						   =						(TFResNetModel, TFResNetForImageClassification) if is_tf_available() else ()
   __UpperCamelCase :			str						   =						(
       {'feature-extraction': TFResNetModel, 'image-classification': TFResNetForImageClassification}
       if is_tf_available()
       else {}
   )

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





   def   UpperCAmelCase__  (       self					:      Optional[int]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							TFResNetModelTester(self     )
      __SCREAMING_SNAKE_CASE   				=							ConfigTester(self					,config_class=lowerCamelCase					,has_text_modality=lowerCamelCase     )





   def   UpperCAmelCase__  (       self					:      Dict     ):

      '''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					:      Dict     ):

      '''simple docstring'''

      return





   @unittest.skip(reason="""ResNet does not use inputs_embeds"""     )
   def   UpperCAmelCase__  (       self					:      Dict     ):

      '''simple docstring'''

      pass





   @unittest.skip(reason="""ResNet does not support input and output embeddings"""     )
   def   UpperCAmelCase__  (       self					:      Optional[int]     ):

      '''simple docstring'''

      pass





   def   UpperCAmelCase__  (       self					:      List[Any]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE   				=							self.model_tester.prepare_config_and_inputs_for_common()

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

         __SCREAMING_SNAKE_CASE   				=							["""pixel_values"""]
         self.assertListEqual(arg_names[:1]					,lowerCamelCase     )





   def   UpperCAmelCase__  (       self					:      Union[str, Any]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*lowerCamelCase     )





   def   UpperCAmelCase__  (       self					:      Dict     ):

      '''simple docstring'''

      def check_hidden_states_output(lowerCamelCase					:      List[str]					,lowerCamelCase					:      Dict					,lowerCamelCase					:      str     ):
         __SCREAMING_SNAKE_CASE   				=							model_class(lowerCamelCase     )
         __SCREAMING_SNAKE_CASE   				=							model(**self._prepare_for_class(lowerCamelCase					,lowerCamelCase     )     )

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

         __SCREAMING_SNAKE_CASE   				=							self.model_tester.num_stages
         self.assertEqual(len(lowerCamelCase     )					,expected_num_stages + 1     )

         # ResNet'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]					,)

      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE   				=							self.model_tester.prepare_config_and_inputs_for_common()
      __SCREAMING_SNAKE_CASE   				=							["""basic""", """bottleneck"""]
      for model_class in self.all_model_classes:
         for layer_type in layers_type:
            __SCREAMING_SNAKE_CASE   				=							layer_type
            __SCREAMING_SNAKE_CASE   				=							True
            check_hidden_states_output(lowerCamelCase					,lowerCamelCase					,lowerCamelCase     )

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

            check_hidden_states_output(lowerCamelCase					,lowerCamelCase					,lowerCamelCase     )





   def   UpperCAmelCase__  (       self					:      str     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*lowerCamelCase     )





   @slow
   def   UpperCAmelCase__  (       self					:      Any     ):

      '''simple docstring'''

      for model_name in TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         __SCREAMING_SNAKE_CASE   				=							TFResNetModel.from_pretrained(lowerCamelCase     )
         self.assertIsNotNone(lowerCamelCase     )





def     __magic_name__							(				)    ->       Optional[Any]:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""   )
   return image







@require_tf
@require_vision
class __a			( unittest.TestCase      ):





   @cached_property
   def   UpperCAmelCase__  (       self					:      Optional[Any]     ):

      '''simple docstring'''

      return (
          AutoImageProcessor.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]     )
          if is_vision_available()
          else None
      )





   @slow
   def   UpperCAmelCase__  (       self					:      Optional[int]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							TFResNetForImageClassification.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]     )

      __SCREAMING_SNAKE_CASE   				=							self.default_image_processor
      __SCREAMING_SNAKE_CASE   				=							prepare_img()
      __SCREAMING_SNAKE_CASE   				=							image_processor(images=lowerCamelCase					,return_tensors="""tf"""     )

      # forward pass
      __SCREAMING_SNAKE_CASE   				=							model(**lowerCamelCase     )

      # verify the logits
      __SCREAMING_SNAKE_CASE   				=							tf.TensorShape((1, 1000)     )
      self.assertEqual(outputs.logits.shape					,lowerCamelCase     )

      __SCREAMING_SNAKE_CASE   				=							tf.constant([-11.1_069, -9.7_877, -8.3_777]     )

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












'''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
a   	=    {
    "/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     __magic_name__							(				__UpperCAmelCase   )    ->       Optional[int]:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							list(s_dict.keys()   )
   for key in keys:
      __SCREAMING_SNAKE_CASE   				=							R""".*/layers_(\d+)"""
      __SCREAMING_SNAKE_CASE   				=							key
      if re.match(__UpperCAmelCase				,						__UpperCAmelCase   ):
         __SCREAMING_SNAKE_CASE   				=							re.sub(R"""layers_(\d+)"""				,						R"""block/\1/layer"""				,						__UpperCAmelCase   )

      __SCREAMING_SNAKE_CASE   				=							R"""(encoder|decoder)\/"""

      if re.match(__UpperCAmelCase				,						__UpperCAmelCase   ):
         __SCREAMING_SNAKE_CASE   				=							re.match(__UpperCAmelCase				,						__UpperCAmelCase   ).groups()
         if groups[0] == "encoder":
            __SCREAMING_SNAKE_CASE   				=							re.sub(R"""/mlp/"""				,						R"""/1/mlp/"""				,						__UpperCAmelCase   )
            __SCREAMING_SNAKE_CASE   				=							re.sub(R"""/pre_mlp_layer_norm/"""				,						R"""/1/layer_norm/"""				,						__UpperCAmelCase   )

         elif groups[0] == "decoder":
            __SCREAMING_SNAKE_CASE   				=							re.sub(R"""/mlp/"""				,						R"""/2/mlp/"""				,						__UpperCAmelCase   )
            __SCREAMING_SNAKE_CASE   				=							re.sub(R"""/pre_mlp_layer_norm/"""				,						R"""/2/layer_norm/"""				,						__UpperCAmelCase   )

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

      print(f"""{key} -> {new_key}"""   )
      __SCREAMING_SNAKE_CASE   				=							s_dict.pop(__UpperCAmelCase   )

   if "encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
      __SCREAMING_SNAKE_CASE   				=							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:
      __SCREAMING_SNAKE_CASE   				=							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:
         __SCREAMING_SNAKE_CASE   				=							s_dict[key].shape[0]
         __SCREAMING_SNAKE_CASE   				=							s_dict[key]
         for idx in range(__UpperCAmelCase   ):
            __SCREAMING_SNAKE_CASE   				=							expert_weihts[idx]
            print(f"""{key} -> {key.replace('expert/'				,						'nested fstring'   )}"""   )

         s_dict.pop(__UpperCAmelCase   )

   return s_dict


a   	=    {
    "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     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase   )    ->       str:

   '''simple docstring'''







   import regex as re

   with open(__UpperCAmelCase				,						"""r"""   ) as f:
      __SCREAMING_SNAKE_CASE   				=							f.read()

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

   __SCREAMING_SNAKE_CASE   				=							re.findall(R"""(.*activations) = \(\'(.*)\',\)"""				,						__UpperCAmelCase   )[0]
   __SCREAMING_SNAKE_CASE   				=							str(activation[1]   )

   __SCREAMING_SNAKE_CASE   				=							num_experts
   __SCREAMING_SNAKE_CASE   				=							SwitchTransformersConfig(**__UpperCAmelCase   )
   return config

def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase				,						__UpperCAmelCase=None				,						__UpperCAmelCase="./"				,						__UpperCAmelCase=8   )    ->       str:

   '''simple docstring'''







   print(f"""Loading flax weights from : {flax_checkpoint_path}"""   )
   __SCREAMING_SNAKE_CASE   				=							checkpoints.load_tax_checkpoint(__UpperCAmelCase   )

   if gin_file is not None:
      __SCREAMING_SNAKE_CASE   				=							convert_gin_to_config(__UpperCAmelCase				,						__UpperCAmelCase   )
   else:
      __SCREAMING_SNAKE_CASE   				=							SwitchTransformersConfig.from_pretrained(__UpperCAmelCase   )

   __SCREAMING_SNAKE_CASE   				=							SwitchTransformersForConditionalGeneration(__UpperCAmelCase   )

   __SCREAMING_SNAKE_CASE   				=							flax_params["""target"""]
   __SCREAMING_SNAKE_CASE   				=							flatten_dict(__UpperCAmelCase				,						sep="""/"""   )
   __SCREAMING_SNAKE_CASE   				=							rename_keys(__UpperCAmelCase   )
   __SCREAMING_SNAKE_CASE   				=							unflatten_dict(__UpperCAmelCase				,						sep="""/"""   )

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

   print(f"""Save PyTorch model to {pytorch_dump_path}"""   )
   pt_model.save_pretrained(__UpperCAmelCase   )


if __name__ == "__main__":
    a   	=    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")
    a   	=    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,
    )





