





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
_lowerCAmelCase				:   Optional[Any]						       =						{
    '''/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   __snake_case    (       _lowerCAmelCase			:			List[str]						)			->    Optional[Any]:
   # 1. in HF T5, we have block.{x}.layer.{y}. which corresponds to layer.{x} in
   # the original model
   A_       :					List[Any]        =      list(s_dict.keys()						)
   for key in keys:
      A_       :					Optional[Any]        =      r".*/layers_(\d+)"
      A_       :					List[Any]        =      key
      if re.match(_lowerCAmelCase					,					_lowerCAmelCase						):
         A_       :					Optional[Any]        =      re.sub(r"layers_(\d+)"					,					r"block/\1/layer"					,					_lowerCAmelCase						)

      A_       :					Union[str, Any]        =      r"(encoder|decoder)\/"

      if re.match(_lowerCAmelCase					,					_lowerCAmelCase						):
         A_       :					Optional[int]        =      re.match(_lowerCAmelCase					,					_lowerCAmelCase						).groups()
         if groups[0] == "encoder":
            A_       :					List[Any]        =      re.sub(r"/mlp/"					,					r"/1/mlp/"					,					_lowerCAmelCase						)
            A_       :					str        =      re.sub(r"/pre_mlp_layer_norm/"					,					r"/1/layer_norm/"					,					_lowerCAmelCase						)

         elif groups[0] == "decoder":
            A_       :					List[Any]        =      re.sub(r"/mlp/"					,					r"/2/mlp/"					,					_lowerCAmelCase						)
            A_       :					Dict        =      re.sub(r"/pre_mlp_layer_norm/"					,					r"/2/layer_norm/"					,					_lowerCAmelCase						)

        # 2. Convert other classic mappings
      for old_key, temp_key in MOE_LAYER_NAME_MAPPING.items():
         if old_key in new_key:
            A_       :					Union[str, Any]        =      new_key.replace(_lowerCAmelCase					,					_lowerCAmelCase						)

      print(f"{key} -> {new_key}"						)
      A_       :					Optional[Any]        =      s_dict.pop(_lowerCAmelCase						)

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

         s_dict.pop(_lowerCAmelCase						)

   return s_dict


_lowerCAmelCase				:   Dict						       =						{
    '''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   __snake_case    (       _lowerCAmelCase			:			Dict					,					_lowerCAmelCase			:			int						)			->    Any:
   # Convert a google style config to the hugging face fromat
   import regex as re

   with open(_lowerCAmelCase					,					"r"						) as f:
      A_       :					Optional[Any]        =      f.read()

   A_       :					Tuple        =      re.findall(r"(.*) = ([0-9.]*)"					,					_lowerCAmelCase						)
   A_       :					Optional[int]        =      {}
   for param, value in regex_match:
      if param in GIN_TO_CONFIG_MAPPING and value != "":
         A_       :					Dict        =      float(_lowerCAmelCase						) if "." in value else int(_lowerCAmelCase						)

   A_       :					Dict        =      re.findall(r"(.*activations) = \(\'(.*)\',\)"					,					_lowerCAmelCase						)[0]
   A_       :					Optional[int]        =      str(activation[1]						)

   A_       :					Union[str, Any]        =      num_experts
   A_       :					str        =      SwitchTransformersConfig(**_lowerCAmelCase						)
   return config

def   __snake_case    (       _lowerCAmelCase			:			Union[str, Any]					,					_lowerCAmelCase			:			List[Any]					,					_lowerCAmelCase			:			str=None					,					_lowerCAmelCase			:			Dict="./"					,					_lowerCAmelCase			:			List[str]=8						)			->    List[Any]:
   # Initialise PyTorch model

   print(f"Loading flax weights from : {flax_checkpoint_path}"						)
   A_       :					int        =      checkpoints.load_tax_checkpoint(_lowerCAmelCase						)

   if gin_file is not None:
      A_       :					int        =      convert_gin_to_config(_lowerCAmelCase					,					_lowerCAmelCase						)
   else:
      A_       :					List[str]        =      SwitchTransformersConfig.from_pretrained(_lowerCAmelCase						)

   A_       :					int        =      SwitchTransformersForConditionalGeneration(_lowerCAmelCase						)

   A_       :					Dict        =      flax_params["target"]
   A_       :					List[Any]        =      flatten_dict(_lowerCAmelCase					,					sep="/"						)
   A_       :					Union[str, Any]        =      rename_keys(_lowerCAmelCase						)
   A_       :					str        =      unflatten_dict(_lowerCAmelCase					,					sep="/"						)

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

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


if __name__ == "__main__":
 _lowerCAmelCase				:   int						       =						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''')
 _lowerCAmelCase				:   Union[str, Any]						       =						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,
 )







# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch

from accelerate import PartialState
from accelerate.utils.operations import broadcast, gather, gather_object, pad_across_processes, reduce

def   __snake_case    (       _lowerCAmelCase			:			Optional[int]						)			->    Dict:
   return (torch.arange(state.num_processes						) + 1.0 + (state.num_processes * state.process_index)).to(state.device						)

def   __snake_case    (       _lowerCAmelCase			:			Tuple						)			->    Optional[Any]:
   A_       :					Optional[Any]        =      create_tensor(_lowerCAmelCase						)
   A_       :					Dict        =      gather(_lowerCAmelCase						)
   assert gathered_tensor.tolist() == list(range(1					,					state.num_processes**2 + 1						)						)

def   __snake_case    (       _lowerCAmelCase			:			List[Any]						)			->    Any:
   A_       :					int        =      [state.process_index]
   A_       :					Union[str, Any]        =      gather_object(_lowerCAmelCase						)
   assert len(_lowerCAmelCase						) == state.num_processes, f"{gathered_obj}, {len(_lowerCAmelCase						)} != {state.num_processes}"
   assert gathered_obj == list(range(state.num_processes						)						), f"{gathered_obj} != {list(range(state.num_processes						)						)}"

def   __snake_case    (       _lowerCAmelCase			:			Union[str, Any]						)			->    Optional[Any]:
   A_       :					List[str]        =      create_tensor(_lowerCAmelCase						)
   A_       :					Optional[Any]        =      broadcast(_lowerCAmelCase						)
   assert broadcasted_tensor.shape == torch.Size([state.num_processes]						)
   assert broadcasted_tensor.tolist() == list(range(1					,					state.num_processes + 1						)						)

def   __snake_case    (       _lowerCAmelCase			:			Dict						)			->    str:
   # We need to pad the tensor with one more element if we are the main process
   # to ensure that we can pad
   if state.is_main_process:
      A_       :					Tuple        =      torch.arange(state.num_processes + 1						).to(state.device						)
   else:
      A_       :					List[str]        =      torch.arange(state.num_processes						).to(state.device						)
   A_       :					Any        =      pad_across_processes(_lowerCAmelCase						)
   assert padded_tensor.shape == torch.Size([state.num_processes + 1]						)
   if not state.is_main_process:
      assert padded_tensor.tolist() == list(range(0					,					state.num_processes						)						) + [0]

def   __snake_case    (       _lowerCAmelCase			:			List[Any]						)			->    Tuple:
   # For now runs on only two processes
   if state.num_processes != 2:
      return
   A_       :					str        =      create_tensor(_lowerCAmelCase						)
   A_       :					int        =      reduce(_lowerCAmelCase					,					"sum"						)
   A_       :					Optional[Any]        =      torch.tensor([4.0, 6]						).to(state.device						)
   assert torch.allclose(_lowerCAmelCase					,					_lowerCAmelCase						), f"{reduced_tensor} != {truth_tensor}"

def   __snake_case    (       _lowerCAmelCase			:			Union[str, Any]						)			->    Dict:
   # For now runs on only two processes
   if state.num_processes != 2:
      return
   A_       :					List[str]        =      create_tensor(_lowerCAmelCase						)
   A_       :					Tuple        =      reduce(_lowerCAmelCase					,					"mean"						)
   A_       :					List[str]        =      torch.tensor([2.0, 3]						).to(state.device						)
   assert torch.allclose(_lowerCAmelCase					,					_lowerCAmelCase						), f"{reduced_tensor} != {truth_tensor}"

def   __snake_case    (       _lowerCAmelCase			:			List[str]						)			->    Dict:
   # For xla_spawn (TPUs)
   main()

def   __snake_case    (       )			->    List[str]:
   A_       :					Tuple        =      PartialState()
   state.print(f"State: {state}"						)
   state.print("testing gather"						)
   test_gather(_lowerCAmelCase						)
   state.print("testing gather_object"						)
   test_gather_object(_lowerCAmelCase						)
   state.print("testing broadcast"						)
   test_broadcast(_lowerCAmelCase						)
   state.print("testing pad_across_processes"						)
   test_pad_across_processes(_lowerCAmelCase						)
   state.print("testing reduce_sum"						)
   test_reduce_sum(_lowerCAmelCase						)
   state.print("testing reduce_mean"						)
   test_reduce_mean(_lowerCAmelCase						)


if __name__ == "__main__":
 main()
