



import argparse
import json
import os

import tensorstore as ts
import torch
from flax import serialization
from flax.traverse_util import flatten_dict, unflatten_dict
from tensorflow.io import gfile

from transformers.modeling_utils import dtype_byte_size
from transformers.models.switch_transformers.convert_switch_transformers_original_flax_checkpoint_to_pytorch import (
    rename_keys,
)
from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME
from transformers.utils.hub import convert_file_size_to_int

def    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_    ) ->       Optional[Any]:
						if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 3:
												# expert layer
												UpperCAmelCase         =     flax_key_tuple[:-1] + ("weight",)
												UpperCAmelCase         =     torch.permute(lowerCamelCase_    ,   (0, 2, 1)    )
						elif flax_key_tuple[-1] == "kernel" and ".".join(lowerCamelCase_    ):
												# linear layer
												UpperCAmelCase         =     flax_key_tuple[:-1] + ("weight",)
												UpperCAmelCase         =     flax_tensor.T
						elif flax_key_tuple[-1] in ["scale", "embedding"]:
												UpperCAmelCase         =     flax_key_tuple[:-1] + ("weight",)

						return flax_key_tuple, flax_tensor

def    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_    ) ->       int:
						if "metadata" in layer:
												UpperCAmelCase         =     layer.split("metadata"    )
												UpperCAmelCase         =     "".join(split_layer[0]    )[:-1]
												UpperCAmelCase         =     [tuple(("metadata" + split_layer[1]).split("/"    )    )]
						elif "kvstore" in layer:
												UpperCAmelCase         =     layer.split("kvstore"    )
												UpperCAmelCase         =     "".join(split_layer[0]    )[:-1]
												UpperCAmelCase         =     [tuple(("kvstore" + split_layer[1]).split("/"    )    )]

						else:
												UpperCAmelCase         =     layer.split("/"    )
												UpperCAmelCase         =     "/".join(split_layer[:-1]    )
												UpperCAmelCase         =     (split_layer[-1],)

						if "kvstore/path" in layer:
												UpperCAmelCase         =     F'{switch_checkpoint_path}/{checkpoint_info[layer]}'
						elif "kvstore/driver" in layer:
												UpperCAmelCase         =     "file"
						else:
												UpperCAmelCase         =     checkpoint_info[layer]

						return curr_real_layer_name, split_layer, content

def    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_    ) ->       List[str]:
						UpperCAmelCase         =     rename_keys(lowerCamelCase_    )
						UpperCAmelCase         =     {}
						for k, v in current_block.items():
												UpperCAmelCase         =     v
						UpperCAmelCase         =     new_current_block
						torch.save(lowerCamelCase_    ,   lowerCamelCase_    )

def    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_ = WEIGHTS_NAME    ) ->       List[str]:
						UpperCAmelCase         =     convert_file_size_to_int(lowerCamelCase_    )
						UpperCAmelCase         =     []
						UpperCAmelCase         =     {}
						UpperCAmelCase         =     0
						UpperCAmelCase         =     0

						os.makedirs(lowerCamelCase_    ,   exist_ok=lowerCamelCase_    )
						with gfile.GFile(switch_checkpoint_path + "/checkpoint"    ,   "rb"    ) as fp:
												UpperCAmelCase         =     serialization.msgpack_restore(fp.read()    )["optimizer"]["target"]
												UpperCAmelCase         =     flatten_dict(lowerCamelCase_    ,   sep="/"    )

						UpperCAmelCase         =     {}
						for layer in checkpoint_info.keys():
												UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase         =     get_key_and_tensorstore_dict(
												    lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_    )
												if curr_real_layer_name in all_layers:
																		UpperCAmelCase         =     content
												else:
																		UpperCAmelCase         =     {split_layer[-1]: content}

						for key in all_layers.keys():
												# open tensorstore file
												UpperCAmelCase         =     ts.open(unflatten_dict(all_layers[key]    )    ).result().read().result()
												UpperCAmelCase         =     torch.tensor(lowerCamelCase_    )
												UpperCAmelCase         =     raw_weights.numel() * dtype_byte_size(raw_weights.dtype    )

												# use the renaming pattern from the small conversion scripts
												UpperCAmelCase       ,		UpperCAmelCase         =     rename_base_flax_keys(tuple(key.split("/"    )    )    ,   lowerCamelCase_    )
												UpperCAmelCase         =     "/".join(lowerCamelCase_    )

												# If this weight is going to tip up over the maximal size, we split.
												if current_block_size + weight_size > max_shard_size:
																		UpperCAmelCase         =     os.path.join(
																		    lowerCamelCase_    ,   weights_name.replace(".bin"    ,   F'-{len(lowerCamelCase_    )+1:05d}-of-???.bin'    )    )
																		rename_and_save_block(lowerCamelCase_    ,   lowerCamelCase_    )
																		sharded_state_dicts.append(current_block.keys()    )
																		del current_block
																		UpperCAmelCase         =     {}
																		UpperCAmelCase         =     0

												UpperCAmelCase         =     raw_weights.to(getattr(lowerCamelCase_    ,   lowerCamelCase_    )    )
												current_block_size += weight_size
												total_size += weight_size

						# Add the last block
						UpperCAmelCase         =     os.path.join(lowerCamelCase_    ,   weights_name.replace(".bin"    ,   F'-{len(lowerCamelCase_    )+1:05d}-of-???.bin'    )    )
						rename_and_save_block(lowerCamelCase_    ,   lowerCamelCase_    )
						sharded_state_dicts.append(current_block.keys()    )

						# If we only have one shard, we return it
						if len(lowerCamelCase_    ) == 1:
												return {weights_name: sharded_state_dicts[0]}, None

						# Otherwise, let's build the index
						UpperCAmelCase         =     {}
						UpperCAmelCase         =     {}
						for idx, shard in enumerate(lowerCamelCase_    ):
												UpperCAmelCase         =     weights_name.replace(
												    ".bin"    ,   F'-{idx+1:05d}-of-{len(lowerCamelCase_    ):05d}.bin'    )  # len(sharded_state_dicts):05d}
												UpperCAmelCase         =     os.path.join(lowerCamelCase_    ,   weights_name.replace(".bin"    ,   F'-{idx+1:05d}-of-???.bin'    )    )
												os.rename(lowerCamelCase_    ,   os.path.join(lowerCamelCase_    ,   lowerCamelCase_    )    )
												UpperCAmelCase         =     shard
												for key in shard:
																		UpperCAmelCase         =     shard_file

    # Add the metadata
						UpperCAmelCase         =     {"total_size": total_size}
						UpperCAmelCase         =     {"metadata": metadata, "weight_map": weight_map}

						with open(os.path.join(lowerCamelCase_    ,   lowerCamelCase_    )    ,   "w"    ,   encoding="utf-8"    ) as f:
												UpperCAmelCase         =     json.dumps(lowerCamelCase_    ,   indent=2    ,   sort_keys=lowerCamelCase_    ) + "\n"
												f.write(lowerCamelCase_    )

						return metadata, index


if __name__ == "__main__":
		__lowerCamelCase     :							Optional[Any]  							=				argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    "--switch_t5x_checkpoint_path",
		    default="/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128/checkpoint_634600",
		    type=str,
		    required=False,
		    help="Path to a directory containing a folder per layer. Follows the original Google format.",
		)
		parser.add_argument("--max_shard_size", default="10GB", required=False, help="Max shard size")
		parser.add_argument("--dtype", default="bfloat16", type=str, required=False, help="dtype of the saved model")
		parser.add_argument(
		    "--pytorch_dump_folder_path",
		    default="/mnt/disks/disk_switch/original_checkpoints/switch-xxl-128-converted",
		    type=str,
		    required=False,
		    help="Path to the output pytorch model.",
		)
		__lowerCamelCase     :							Union[str, Any]  							=				parser.parse_args()
		shard_on_the_fly(
		    args.switch_tax_checkpoint_path,
		    args.pytorch_dump_folder_path,
		    args.max_shard_size,
		    args.dtype,
		)





def    lowerCamelCase_() ->       Dict:
						from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration, TaTokenizer

						UpperCAmelCase         =     SwitchTransformersConfig.from_pretrained("google/switch-base-8"    )
						config.save_pretrained("/home/arthur_huggingface_co/transformers/switch_converted"    )
						UpperCAmelCase         =     SwitchTransformersForConditionalGeneration.from_pretrained(
						    "/home/arthur_huggingface_co/transformers/switch_converted"    ,   device_map="auto"    )

						UpperCAmelCase         =     TaTokenizer.from_pretrained("t5-small"    )
						UpperCAmelCase         =     "A <extra_id_0> walks into a bar a orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>."

						UpperCAmelCase         =     tokenizer(lowerCamelCase_    ,   return_tensors="pt"    ).input_ids
						UpperCAmelCase         =     model.generate(lowerCamelCase_    ,   decoder_start_token_id=0    )
						print(tokenizer.decode(out[0]    )    )








import os
import re
import shutil
import sys
import tempfile
import unittest

import black


__lowerCamelCase     :							str  							=				os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, "utils"))

import check_copies  # noqa: E402


# This is the reference code that will be used in the tests.
# If BertLMPredictionHead is changed in modeling_bert.py, this code needs to be manually updated.
__lowerCamelCase     :							Optional[Any]  							=				"    def __init__(self, config):\n        super().__init__()\n        self.transform = BertPredictionHeadTransform(config)\n\n        # The output weights are the same as the input embeddings, but there is\n        # an output-only bias for each token.\n        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)\n\n        self.bias = nn.Parameter(torch.zeros(config.vocab_size))\n\n        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`\n        self.decoder.bias = self.bias\n\n    def forward(self, hidden_states):\n        hidden_states = self.transform(hidden_states)\n        hidden_states = self.decoder(hidden_states)\n        return hidden_states\n"





class  __magic_name__   (							unittest.TestCase						):






		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[int]      )				->					Optional[Any]:
								'''simple docstring'''







								UpperCAmelCase         =     tempfile.mkdtemp()
								os.makedirs(os.path.join(self.transformer_dir				,		"models/bert/"      )      )
								UpperCAmelCase         =     self.transformer_dir
								shutil.copy(
								    os.path.join(UpperCamelCase__				,		"src/transformers/models/bert/modeling_bert.py"      )				,		os.path.join(self.transformer_dir				,		"models/bert/modeling_bert.py"      )				,		)






		def SCREAMING_SNAKE_CASE_			(						self      :     int      )				->					List[Any]:
								'''simple docstring'''







								UpperCAmelCase         =     "src/transformers"
								shutil.rmtree(self.transformer_dir      )






		def SCREAMING_SNAKE_CASE_			(						self      :     Tuple				,		UpperCamelCase__      :     List[Any]				,		UpperCamelCase__      :     Optional[int]				,		UpperCamelCase__      :     Tuple				,		UpperCamelCase__      :     Union[str, Any]=None      )				->					str:
								'''simple docstring'''







								UpperCAmelCase         =     comment + F'\nclass {class_name}(nn.Module):\n' + class_code
								if overwrite_result is not None:
														UpperCAmelCase         =     comment + F'\nclass {class_name}(nn.Module):\n' + overwrite_result
								UpperCAmelCase         =     black.Mode(target_versions={black.TargetVersion.PYaa}				,		line_length=1_19      )
								UpperCAmelCase         =     black.format_str(UpperCamelCase__				,		mode=UpperCamelCase__      )
								UpperCAmelCase         =     os.path.join(self.transformer_dir				,		"new_code.py"      )
								with open(UpperCamelCase__				,		"w"				,		newline="\n"      ) as f:
														f.write(UpperCamelCase__      )
								if overwrite_result is None:
														self.assertTrue(len(check_copies.is_copy_consistent(UpperCamelCase__      )      ) == 0      )
								else:
														check_copies.is_copy_consistent(f.name				,		overwrite=UpperCamelCase__      )
														with open(UpperCamelCase__				,		"r"      ) as f:
																				self.assertTrue(f.read()				,		UpperCamelCase__      )






		def SCREAMING_SNAKE_CASE_			(						self      :     List[Any]      )				->					Optional[Any]:
								'''simple docstring'''







								UpperCAmelCase         =     check_copies.find_code_in_transformers("models.bert.modeling_bert.BertLMPredictionHead"      )
								self.assertEqual(UpperCamelCase__				,		UpperCamelCase__      )






		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[Any]      )				->					Dict:
								'''simple docstring'''







								self.check_copy_consistency(
								    "# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead"				,		"BertLMPredictionHead"				,		REFERENCE_CODE + "\n"				,		)

								# With no empty line at the end
								self.check_copy_consistency(
								    "# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead"				,		"BertLMPredictionHead"				,		UpperCamelCase__				,		)

								# Copy consistency with rename
								self.check_copy_consistency(
								    "# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead with Bert->TestModel"				,		"TestModelLMPredictionHead"				,		re.sub("Bert"				,		"TestModel"				,		UpperCamelCase__      )				,		)

								# Copy consistency with a really long name
								UpperCAmelCase         =     "TestModelWithAReallyLongNameBecauseSomePeopleLikeThatForSomeReason"
								self.check_copy_consistency(
								    F'# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead with Bert->{long_class_name}'				,		F'{long_class_name}LMPredictionHead'				,		re.sub("Bert"				,		UpperCamelCase__				,		UpperCamelCase__      )				,		)

								# Copy consistency with overwrite
								self.check_copy_consistency(
								    "# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead with Bert->TestModel"				,		"TestModelLMPredictionHead"				,		UpperCamelCase__				,		overwrite_result=re.sub("Bert"				,		"TestModel"				,		UpperCamelCase__      )				,		)







		def SCREAMING_SNAKE_CASE_			(						self      :     int      )				->					Optional[Any]:
								'''simple docstring'''







								UpperCAmelCase         =     check_copies.LOCALIZED_READMES["README_zh-hans.md"]

								UpperCAmelCase         =     (
								    "1. **[ALBERT](https://huggingface.co/transformers/model_doc/albert.html)** (from Google Research and the"
								    " Toyota Technological Institute at Chicago) released with the paper [ALBERT: A Lite BERT for"
								    " Self-supervised Learning of Language Representations](https://arxiv.org/abs/1909.11942), by Zhenzhong"
								    " Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut.\n1."
								    " **[DistilBERT](https://huggingface.co/transformers/model_doc/distilbert.html)** (from HuggingFace),"
								    " released together with the paper [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and"
								    " lighter](https://arxiv.org/abs/1910.01108) by Victor Sanh, Lysandre Debut and Thomas Wolf. The same"
								    " method has been applied to compress GPT2 into"
								    " [DistilGPT2](https://github.com/huggingface/transformers/tree/main/examples/distillation), RoBERTa into"
								    " [DistilRoBERTa](https://github.com/huggingface/transformers/tree/main/examples/distillation),"
								    " Multilingual BERT into"
								    " [DistilmBERT](https://github.com/huggingface/transformers/tree/main/examples/distillation) and a German"
								    " version of DistilBERT.\n1. **[ELECTRA](https://huggingface.co/transformers/model_doc/electra.html)**"
								    " (from Google Research/Stanford University) released with the paper [ELECTRA: Pre-training text encoders"
								    " as discriminators rather than generators](https://arxiv.org/abs/2003.10555) by Kevin Clark, Minh-Thang"
								    " Luong, Quoc V. Le, Christopher D. Manning."
								)
								UpperCAmelCase         =     (
								    "1. **[ALBERT](https://huggingface.co/transformers/model_doc/albert.html)** (来自 Google Research and the"
								    " Toyota Technological Institute at Chicago) 伴随论文 [ALBERT: A Lite BERT for Self-supervised Learning of"
								    " Language Representations](https://arxiv.org/abs/1909.11942), 由 Zhenzhong Lan, Mingda Chen, Sebastian"
								    " Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut 发布。\n"
								)
								UpperCAmelCase         =     (
								    "1. **[ALBERT](https://huggingface.co/transformers/model_doc/albert.html)** (来自 Google Research and the"
								    " Toyota Technological Institute at Chicago) 伴随论文 [ALBERT: A Lite BERT for Self-supervised Learning of"
								    " Language Representations](https://arxiv.org/abs/1909.11942), 由 Zhenzhong Lan, Mingda Chen, Sebastian"
								    " Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut 发布。\n1."
								    " **[DistilBERT](https://huggingface.co/transformers/model_doc/distilbert.html)** (来自 HuggingFace) 伴随论文"
								    " [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and"
								    " lighter](https://arxiv.org/abs/1910.01108) 由 Victor Sanh, Lysandre Debut and Thomas Wolf 发布。 The same"
								    " method has been applied to compress GPT2 into"
								    " [DistilGPT2](https://github.com/huggingface/transformers/tree/main/examples/distillation), RoBERTa into"
								    " [DistilRoBERTa](https://github.com/huggingface/transformers/tree/main/examples/distillation),"
								    " Multilingual BERT into"
								    " [DistilmBERT](https://github.com/huggingface/transformers/tree/main/examples/distillation) and a German"
								    " version of DistilBERT.\n1. **[ELECTRA](https://huggingface.co/transformers/model_doc/electra.html)** (来自"
								    " Google Research/Stanford University) 伴随论文 [ELECTRA: Pre-training text encoders as discriminators rather"
								    " than generators](https://arxiv.org/abs/2003.10555) 由 Kevin Clark, Minh-Thang Luong, Quoc V. Le,"
								    " Christopher D. Manning 发布。\n"
								)

								UpperCAmelCase       ,		UpperCAmelCase         =     check_copies.convert_to_localized_md(
								    UpperCamelCase__				,		UpperCamelCase__				,		localized_readme["format_model_list"]      )

								self.assertFalse(UpperCamelCase__      )
								self.assertEqual(UpperCamelCase__				,		UpperCamelCase__      )

								UpperCAmelCase       ,		UpperCAmelCase         =     check_copies.convert_to_localized_md(
								    UpperCamelCase__				,		UpperCamelCase__				,		localized_readme["format_model_list"]      )

								# Check whether the number of models is equal to README.md after conversion.
								self.assertTrue(UpperCamelCase__      )

								UpperCAmelCase         =     (
								    "1. **[ALBERT](https://huggingface.co/transformers/model_doc/albert.html)** (from Google Research and the"
								    " Toyota Technological Institute at Chicago) released with the paper [ALBERT: A Lite BERT for"
								    " Self-supervised Learning of Language Representations](https://arxiv.org/abs/1909.11942), by Zhenzhong"
								    " Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut."
								)
								UpperCAmelCase         =     (
								    "1. **[ALBERT](https://huggingface.co/transformers/main/model_doc/albert.html)** (来自 Google Research and"
								    " the Toyota Technological Institute at Chicago) 伴随论文 [ALBERT: A Lite BERT for Self-supervised Learning of"
								    " Language Representations](https://arxiv.org/abs/1909.11942), 由 Zhenzhong Lan, Mingda Chen, Sebastian"
								    " Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut 发布。\n"
								)
								UpperCAmelCase         =     (
								    "1. **[ALBERT](https://huggingface.co/transformers/model_doc/albert.html)** (来自 Google Research and the"
								    " Toyota Technological Institute at Chicago) 伴随论文 [ALBERT: A Lite BERT for Self-supervised Learning of"
								    " Language Representations](https://arxiv.org/abs/1909.11942), 由 Zhenzhong Lan, Mingda Chen, Sebastian"
								    " Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut 发布。\n"
								)

								UpperCAmelCase       ,		UpperCAmelCase         =     check_copies.convert_to_localized_md(
								    UpperCamelCase__				,		UpperCamelCase__				,		localized_readme["format_model_list"]      )

								# Check if the model link is synchronized.
								self.assertEqual(UpperCamelCase__				,		UpperCamelCase__      )



