


import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


SCREAMING_SNAKE_CASE						:Optional[Any]			=  datasets.utils.logging.get_logger(__name__)


@dataclass
class      __lowerCAmelCase  (			datasets.BuilderConfig       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							None
  _SCREAMING_SNAKE_CASE		     =							"utf-8"
  _SCREAMING_SNAKE_CASE		     =							None
  _SCREAMING_SNAKE_CASE		     =							None
  _SCREAMING_SNAKE_CASE		     =							True  # deprecated
  _SCREAMING_SNAKE_CASE		     =							None  # deprecated
  _SCREAMING_SNAKE_CASE		     =							10 << 20  # 10MB
  _SCREAMING_SNAKE_CASE		     =							None






class      __lowerCAmelCase  (			datasets.ArrowBasedBuilder       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							JsonConfig



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




     """simple docstring"""
     if self.config.block_size is not None:
        logger.warning("The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead"	)
        snake_case_            =       self.config.block_size
     if self.config.use_threads is not True:
        logger.warning(
            "The JSON loader parameter `use_threads` is deprecated and doesn't have any effect anymore."	)
     if self.config.newlines_in_values is not None:
        raise ValueError("The JSON loader parameter `newlines_in_values` is no longer supported"	)
     return datasets.DatasetInfo(features=self.config.features	)



  def        lowerCAmelCase__	(  self							:							Tuple     ,					_lowerCAmelCase							:							Union[str, Any]	) ->			Union[str, Any]:




     """simple docstring"""
     if not self.config.data_files:
        raise ValueError(F'''At least one data file must be specified, but got data_files={self.config.data_files}'''	)
     snake_case_            =       dl_manager.download_and_extract(self.config.data_files	)
     if isinstance(_lowerCAmelCase     ,					(str, list, tuple)	):
        snake_case_            =       data_files
        if isinstance(_lowerCAmelCase     ,					_lowerCAmelCase	):
           snake_case_            =       [files]
        snake_case_            =       [dl_manager.iter_files(_lowerCAmelCase	) for file in files]
        return [datasets.SplitGenerator(name=datasets.Split.TRAIN     ,					gen_kwargs={"files": files}	)]
     snake_case_            =       []
     for split_name, files in data_files.items():
        if isinstance(_lowerCAmelCase     ,					_lowerCAmelCase	):
           snake_case_            =       [files]
        snake_case_            =       [dl_manager.iter_files(_lowerCAmelCase	) for file in files]
        splits.append(datasets.SplitGenerator(name=_lowerCAmelCase     ,					gen_kwargs={"files": files}	)	)
     return splits



  def        lowerCAmelCase__	(  self							:							Optional[int]     ,					_lowerCAmelCase							:							pa.Table	) ->			pa.Table:




     """simple docstring"""
     if self.config.features is not None:
        # adding missing columns
        for column_name in set(self.config.features	) - set(pa_table.column_names	):
           snake_case_            =       self.config.features.arrow_schema.field(_lowerCAmelCase	).type
           snake_case_            =       pa_table.append_column(_lowerCAmelCase     ,					pa.array([None] * len(_lowerCAmelCase	)     ,					type=_lowerCAmelCase	)	)
        # more expensive cast to support nested structures with keys in a different order
        # allows str <-> int/float or str to Audio for example
        snake_case_            =       table_cast(_lowerCAmelCase     ,					self.config.features.arrow_schema	)
     return pa_table



  def        lowerCAmelCase__	(  self							:							int     ,					_lowerCAmelCase							:							Optional[int]	) ->			Tuple:




     """simple docstring"""
     for file_idx, file in enumerate(itertools.chain.from_iterable(_lowerCAmelCase	)	):
        # If the file is one json object and if we need to look at the list of items in one specific field
        if self.config.field is not None:
           with open(_lowerCAmelCase     ,					encoding=self.config.encoding     ,					errors=self.config.encoding_errors	) as f:
              snake_case_            =       json.load(_lowerCAmelCase	)

           # We keep only the field we are interested in
           snake_case_            =       dataset[self.config.field]

           # We accept two format: a list of dicts or a dict of lists
           if isinstance(_lowerCAmelCase     ,					(list, tuple)	):
              snake_case_            =       set().union(*[row.keys() for row in dataset]	)
              snake_case_            =       {col: [row.get(_lowerCAmelCase	) for row in dataset] for col in keys}
           else:
              snake_case_            =       dataset
           snake_case_            =       pa.Table.from_pydict(_lowerCAmelCase	)
           yield file_idx, self._cast_table(_lowerCAmelCase	)

        # If the file has one json object per line
        else:
           with open(_lowerCAmelCase     ,					"rb"	) as f:
              snake_case_            =       0
              # Use block_size equal to the chunk size divided by 32 to leverage multithreading
              # Set a default minimum value of 16kB if the chunk size is really small
              snake_case_            =       max(self.config.chunksize // 3_2     ,					1_6 << 1_0	)
              snake_case_            =       (
                  self.config.encoding_errors if self.config.encoding_errors is not None else "strict"
              )
              while True:
                 snake_case_            =       f.read(self.config.chunksize	)
                 if not batch:
                    break
                 # Finish current line
                 try:
                    batch += f.readline()
                 except (AttributeError, io.UnsupportedOperation):
                    batch += readline(_lowerCAmelCase	)
                 # PyArrow only accepts utf-8 encoded bytes
                 if self.config.encoding != "utf-8":
                    snake_case_            =       batch.decode(self.config.encoding     ,					errors=_lowerCAmelCase	).encode("utf-8"	)
                 try:
                    while True:
                       try:
                          snake_case_            =       paj.read_json(
                              io.BytesIO(_lowerCAmelCase	)     ,					read_options=paj.ReadOptions(block_size=_lowerCAmelCase	)	)
                          break
                       except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
                          if (
                              isinstance(_lowerCAmelCase     ,					pa.ArrowInvalid	)
                              and "straddling" not in str(_lowerCAmelCase	)
                              or block_size > len(_lowerCAmelCase	)
                          ):
                             raise
                          else:
                             # Increase the block size in case it was too small.
                             # The block size will be reset for the next file.
                             logger.debug(
                                 F'''Batch of {len(_lowerCAmelCase	)} bytes couldn\'t be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}.'''	)
                             block_size *= 2
                 except pa.ArrowInvalid as e:
                    try:
                       with open(
                           _lowerCAmelCase     ,					encoding=self.config.encoding     ,					errors=self.config.encoding_errors	) as f:
                          snake_case_            =       json.load(_lowerCAmelCase	)
                    except json.JSONDecodeError:
                       logger.error(F'''Failed to read file \'{file}\' with error {type(_lowerCAmelCase	)}: {e}'''	)
                       raise e
                    # If possible, parse the file as a list of json objects and exit the loop
                    if isinstance(_lowerCAmelCase     ,					_lowerCAmelCase	):  # list is the only sequence type supported in JSON
                       try:
                          snake_case_            =       set().union(*[row.keys() for row in dataset]	)
                          snake_case_            =       {col: [row.get(_lowerCAmelCase	) for row in dataset] for col in keys}
                          snake_case_            =       pa.Table.from_pydict(_lowerCAmelCase	)
                       except (pa.ArrowInvalid, AttributeError) as e:
                          logger.error(F'''Failed to read file \'{file}\' with error {type(_lowerCAmelCase	)}: {e}'''	)
                          raise ValueError(F'''Not able to read records in the JSON file at {file}.'''	) from None
                       yield file_idx, self._cast_table(_lowerCAmelCase	)
                       break
                    else:
                       logger.error(F'''Failed to read file \'{file}\' with error {type(_lowerCAmelCase	)}: {e}'''	)
                       raise ValueError(
                           F'''Not able to read records in the JSON file at {file}. '''
                           F'''You should probably indicate the field of the JSON file containing your records. '''
                           F'''This JSON file contain the following fields: {str(list(dataset.keys()	)	)}. '''
                           F'''Select the correct one and provide it as `field=\'XXX\'` to the dataset loading method. '''	) from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                 yield (file_idx, batch_idx), self._cast_table(_lowerCAmelCase	)
                 batch_idx += 1








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,
   )




