


import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoFeatureExtractor, WavaVecaFeatureExtractor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / '''utils'''))

from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


SCREAMING_SNAKE_CASE						:Dict			=  get_tests_dir('''fixtures''')


class      __lowerCAmelCase  (			unittest.TestCase       ):






  """simple docstring"""





  def        lowerCAmelCase__	(  self							:							List[str]	) ->			Dict:




     """simple docstring"""
     # A mock response for an HTTP head request to emulate server down
     snake_case_            =       mock.Mock()
     snake_case_            =       5_0_0
     snake_case_            =       {}
     snake_case_            =       HTTPError
     snake_case_            =       {}

     # Download this model to make sure it's in the cache.
     snake_case_            =       WavaVecaFeatureExtractor.from_pretrained("hf-internal-testing/tiny-random-wav2vec2"	)
     # Under the mock environment we get a 500 error when trying to reach the model.
     with mock.patch("requests.Session.request"     ,					return_value=_lowerCAmelCase	) as mock_head:
        snake_case_            =       WavaVecaFeatureExtractor.from_pretrained("hf-internal-testing/tiny-random-wav2vec2"	)
        # This check we did call the fake head request
        mock_head.assert_called()



  def        lowerCAmelCase__	(  self							:							Tuple	) ->			str:




     """simple docstring"""
     # This test is for deprecated behavior and can be removed in v5
     snake_case_            =       WavaVecaFeatureExtractor.from_pretrained(
         "https://huggingface.co/hf-internal-testing/tiny-random-wav2vec2/resolve/main/preprocessor_config.json"	)






@is_staging_test
class      __lowerCAmelCase  (			unittest.TestCase       ):






  """simple docstring"""





  @classmethod
  def        lowerCAmelCase__	(  cls							:							Optional[Any]	) ->			Optional[int]:




     """simple docstring"""
     snake_case_            =       TOKEN
     HfFolder.save_token(_lowerCAmelCase	)



  @classmethod
  def        lowerCAmelCase__	(  cls							:							Any	) ->			Any:




     """simple docstring"""
     try:
        delete_repo(token=cls._token     ,					repo_id="test-feature-extractor"	)
     except HTTPError:
        pass

     try:
        delete_repo(token=cls._token     ,					repo_id="valid_org/test-feature-extractor-org"	)
     except HTTPError:
        pass

     try:
        delete_repo(token=cls._token     ,					repo_id="test-dynamic-feature-extractor"	)
     except HTTPError:
        pass



  def        lowerCAmelCase__	(  self							:							Optional[Any]	) ->			int:




     """simple docstring"""
     snake_case_            =       WavaVecaFeatureExtractor.from_pretrained(_lowerCAmelCase	)
     feature_extractor.push_to_hub("test-feature-extractor"     ,					use_auth_token=self._token	)

     snake_case_            =       WavaVecaFeatureExtractor.from_pretrained(F'''{USER}/test-feature-extractor'''	)
     for k, v in feature_extractor.__dict__.items():
        self.assertEqual(_lowerCAmelCase     ,					getattr(_lowerCAmelCase     ,					_lowerCAmelCase	)	)

     # Reset repo
     delete_repo(token=self._token     ,					repo_id="test-feature-extractor"	)

     # Push to hub via save_pretrained
     with tempfile.TemporaryDirectory() as tmp_dir:
        feature_extractor.save_pretrained(
            _lowerCAmelCase     ,					repo_id="test-feature-extractor"     ,					push_to_hub=_lowerCAmelCase     ,					use_auth_token=self._token	)

     snake_case_            =       WavaVecaFeatureExtractor.from_pretrained(F'''{USER}/test-feature-extractor'''	)
     for k, v in feature_extractor.__dict__.items():
        self.assertEqual(_lowerCAmelCase     ,					getattr(_lowerCAmelCase     ,					_lowerCAmelCase	)	)



  def        lowerCAmelCase__	(  self							:							Any	) ->			Any:




     """simple docstring"""
     snake_case_            =       WavaVecaFeatureExtractor.from_pretrained(_lowerCAmelCase	)
     feature_extractor.push_to_hub("valid_org/test-feature-extractor"     ,					use_auth_token=self._token	)

     snake_case_            =       WavaVecaFeatureExtractor.from_pretrained("valid_org/test-feature-extractor"	)
     for k, v in feature_extractor.__dict__.items():
        self.assertEqual(_lowerCAmelCase     ,					getattr(_lowerCAmelCase     ,					_lowerCAmelCase	)	)

     # Reset repo
     delete_repo(token=self._token     ,					repo_id="valid_org/test-feature-extractor"	)

     # Push to hub via save_pretrained
     with tempfile.TemporaryDirectory() as tmp_dir:
        feature_extractor.save_pretrained(
            _lowerCAmelCase     ,					repo_id="valid_org/test-feature-extractor-org"     ,					push_to_hub=_lowerCAmelCase     ,					use_auth_token=self._token	)

     snake_case_            =       WavaVecaFeatureExtractor.from_pretrained("valid_org/test-feature-extractor-org"	)
     for k, v in feature_extractor.__dict__.items():
        self.assertEqual(_lowerCAmelCase     ,					getattr(_lowerCAmelCase     ,					_lowerCAmelCase	)	)



  def        lowerCAmelCase__	(  self							:							List[str]	) ->			List[Any]:




     """simple docstring"""
     CustomFeatureExtractor.register_for_auto_class()
     snake_case_            =       CustomFeatureExtractor.from_pretrained(_lowerCAmelCase	)

     feature_extractor.push_to_hub("test-dynamic-feature-extractor"     ,					use_auth_token=self._token	)

     # This has added the proper auto_map field to the config
     self.assertDictEqual(
         feature_extractor.auto_map     ,					{"AutoFeatureExtractor": "custom_feature_extraction.CustomFeatureExtractor"}     ,					)

     snake_case_            =       AutoFeatureExtractor.from_pretrained(
         F'''{USER}/test-dynamic-feature-extractor'''     ,					trust_remote_code=_lowerCAmelCase	)
     # Can't make an isinstance check because the new_feature_extractor is from the CustomFeatureExtractor class of a dynamic module
     self.assertEqual(new_feature_extractor.__class__.__name__     ,					"CustomFeatureExtractor"	)








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
SCREAMING_SNAKE_CASE						:Tuple			=  {
    '''/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    (				lowerCAmelCase_				:Any				)->Any:
   '''simple docstring'''






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

      snake_case_            =       r"(encoder|decoder)\/"

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

         elif groups[0] == "decoder":
            snake_case_            =       re.sub(r"/mlp/"		,    r"/2/mlp/"		,    lowerCAmelCase_				)
            snake_case_            =       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:
            snake_case_            =       new_key.replace(lowerCAmelCase_		,    lowerCAmelCase_				)

      print(F'''{key} -> {new_key}'''				)
      snake_case_            =       s_dict.pop(lowerCAmelCase_				)

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

         s_dict.pop(lowerCAmelCase_				)

   return s_dict


SCREAMING_SNAKE_CASE						:Optional[Any]			=  {
    '''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    (				lowerCAmelCase_				:List[str]		,    lowerCAmelCase_				:Any				)->Optional[int]:
   '''simple docstring'''






   import regex as re

   with open(lowerCAmelCase_		,    "r"				) as f:
      snake_case_            =       f.read()

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

   snake_case_            =       re.findall(r"(.*activations) = \(\'(.*)\',\)"		,    lowerCAmelCase_				)[0]
   snake_case_            =       str(activation[1]				)

   snake_case_            =       num_experts
   snake_case_            =       SwitchTransformersConfig(**lowerCAmelCase_				)
   return config



def       _lowerCAmelCase    (				lowerCAmelCase_				:Optional[Any]		,    lowerCAmelCase_				:List[Any]		,    lowerCAmelCase_				:Dict=None		,    lowerCAmelCase_				:str="./"		,    lowerCAmelCase_				:Optional[int]=8				)->List[str]:
   '''simple docstring'''






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

   if gin_file is not None:
      snake_case_            =       convert_gin_to_config(lowerCAmelCase_		,    lowerCAmelCase_				)
   else:
      snake_case_            =       SwitchTransformersConfig.from_pretrained(lowerCAmelCase_				)

   snake_case_            =       SwitchTransformersForConditionalGeneration(lowerCAmelCase_				)

   snake_case_            =       flax_params["target"]
   snake_case_            =       flatten_dict(lowerCAmelCase_		,    sep="/"				)
   snake_case_            =       rename_keys(lowerCAmelCase_				)
   snake_case_            =       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__":
   SCREAMING_SNAKE_CASE						: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''')
   SCREAMING_SNAKE_CASE						:int			=  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,
   )




