'''simple docstring'''






import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers.models.esm.modeling_esmfold import EsmForProteinFolding







class lowercase_       :







   """simple docstring"""

   def __init__(      self      :			List[str], UpperCamelCase__      :			Dict, UpperCamelCase__      :			Dict=13, UpperCamelCase__      :			Union[str, Any]=7, UpperCamelCase__      :			List[str]=False, UpperCamelCase__      :			List[Any]=True, UpperCamelCase__      :			str=False, UpperCamelCase__      :			Tuple=False, UpperCamelCase__      :			str=19, UpperCamelCase__      :			Tuple=32, UpperCamelCase__      :			Optional[Any]=5, UpperCamelCase__      :			int=4, UpperCamelCase__      :			str=37, UpperCamelCase__      :			List[str]="gelu", UpperCamelCase__      :			Union[str, Any]=0.1, UpperCamelCase__      :			List[Any]=0.1, UpperCamelCase__      :			int=5_12, UpperCamelCase__      :			int=16, UpperCamelCase__      :			str=2, UpperCamelCase__      :			Any=0.02, UpperCamelCase__      :			Any=3, UpperCamelCase__      :			Any=4, UpperCamelCase__      :			Union[str, Any]=None, )  ->			Tuple:
        _A    		=  parent
        _A    		=  batch_size
        _A    		=  seq_length
        _A    		=  is_training
        _A    		=  use_input_mask
        _A    		=  use_token_type_ids
        _A    		=  use_labels
        _A    		=  vocab_size
        _A    		=  hidden_size
        _A    		=  num_hidden_layers
        _A    		=  num_attention_heads
        _A    		=  intermediate_size
        _A    		=  hidden_act
        _A    		=  hidden_dropout_prob
        _A    		=  attention_probs_dropout_prob
        _A    		=  max_position_embeddings
        _A    		=  type_vocab_size
        _A    		=  type_sequence_label_size
        _A    		=  initializer_range
        _A    		=  num_labels
        _A    		=  num_choices
        _A    		=  scope
   def 	__UpperCAmelCase   (      self      :			Tuple     )  ->			List[Any]:
        _A    		=  ids_tensor([self.batch_size, self.seq_length], self.vocab_size     )

        _A    		=  None
        if self.use_input_mask:
             _A    		=  random_attention_mask([self.batch_size, self.seq_length]     )

        _A    		=  None
        _A    		=  None
        _A    		=  None
        if self.use_labels:
             _A    		=  ids_tensor([self.batch_size], self.type_sequence_label_size     )
             _A    		=  ids_tensor([self.batch_size, self.seq_length], self.num_labels     )
             _A    		=  ids_tensor([self.batch_size], self.num_choices     )

        _A    		=  self.get_config()

        return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
   def 	__UpperCAmelCase   (      self      :			str     )  ->			Dict:
        _A    		=  EsmConfig(
            vocab_size=33, hidden_size=self.hidden_size, pad_token_id=1, num_hidden_layers=self.num_hidden_layers, num_attention_heads=self.num_attention_heads, intermediate_size=self.intermediate_size, hidden_act=self.hidden_act, hidden_dropout_prob=self.hidden_dropout_prob, attention_probs_dropout_prob=self.attention_probs_dropout_prob, max_position_embeddings=self.max_position_embeddings, type_vocab_size=self.type_vocab_size, initializer_range=self.initializer_range, is_folding_model=UpperCamelCase__, esmfold_config={'trunk': {'num_blocks': 2}, 'fp16_esm': False}, )
        return config
   def 	__UpperCAmelCase   (      self      :			Dict, UpperCamelCase__      :			Union[str, Any], UpperCamelCase__      :			Optional[int], UpperCamelCase__      :			str, UpperCamelCase__      :			Dict, UpperCamelCase__      :			List[Any], UpperCamelCase__      :			int     )  ->			List[str]:
        _A    		=  EsmForProteinFolding(config=UpperCamelCase__     ).float()
        model.to(UpperCamelCase__     )
        model.eval()
        _A    		=  model(UpperCamelCase__, attention_mask=UpperCamelCase__     )
        _A    		=  model(UpperCamelCase__     )
        _A    		=  model(UpperCamelCase__     )

        self.parent.assertEqual(result.positions.shape, (8, self.batch_size, self.seq_length, 14, 3)     )
        self.parent.assertEqual(result.angles.shape, (8, self.batch_size, self.seq_length, 7, 2)     )

   def 	__UpperCAmelCase   (      self      :			Any     )  ->			Optional[int]:
        _A    		=  self.prepare_config_and_inputs()
        (
            (
            _A
        )		,       (
            _A
        )		,       (
            _A
        )		,       (
            _A
        )		,       (
            _A
        )		,       (
            _A
        )		,       
        )    		=  config_and_inputs
        _A    		=  {'input_ids': input_ids, 'attention_mask': input_mask}
        return config, inputs_dict







@require_torch
class lowercase_       (		_UpperCamelCase    ,    _UpperCamelCase    ,    unittest.TestCase			):







   """simple docstring"""

   __lowerCAmelCase         =   False

   __lowerCAmelCase         =   (EsmForProteinFolding,) if is_torch_available() else ()
   __lowerCAmelCase         =   ()
   __lowerCAmelCase         =   {} if is_torch_available() else {}
   __lowerCAmelCase         =   False
   def 	__UpperCAmelCase   (      self      :			List[str]     )  ->			Dict:
        _A    		=  EsmFoldModelTester(self     )
        _A    		=  ConfigTester(self, config_class=UpperCamelCase__, hidden_size=37     )
   def 	__UpperCAmelCase   (      self      :			List[Any]     )  ->			Optional[Any]:
        self.config_tester.run_common_tests()
   def 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			str:
        _A    		=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*UpperCamelCase__     )
   @unittest.skip('Does not support attention outputs'     )
   def 	__UpperCAmelCase   (      self      :			Optional[Any]     )  ->			Tuple:
        pass
   @unittest.skip
   def 	__UpperCAmelCase   (      self      :			Tuple     )  ->			Union[str, Any]:
        pass
   @unittest.skip('Esm does not support embedding resizing'     )
   def 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			Any:
        pass
   @unittest.skip('Esm does not support embedding resizing'     )
   def 	__UpperCAmelCase   (      self      :			str     )  ->			Any:
        pass
   @unittest.skip('ESMFold does not support passing input embeds!'     )
   def 	__UpperCAmelCase   (      self      :			Tuple     )  ->			Tuple:
        pass
   @unittest.skip('ESMFold does not support head pruning.'     )
   def 	__UpperCAmelCase   (      self      :			str     )  ->			Optional[Any]:
        pass
   @unittest.skip('ESMFold does not support head pruning.'     )
   def 	__UpperCAmelCase   (      self      :			Any     )  ->			Union[str, Any]:
        pass
   @unittest.skip('ESMFold does not support head pruning.'     )
   def 	__UpperCAmelCase   (      self      :			str     )  ->			Optional[int]:
        pass
   @unittest.skip('ESMFold does not support head pruning.'     )
   def 	__UpperCAmelCase   (      self      :			Dict     )  ->			int:
        pass
   @unittest.skip('ESMFold does not support head pruning.'     )
   def 	__UpperCAmelCase   (      self      :			Dict     )  ->			Optional[int]:
        pass
   @unittest.skip('ESMFold does not output hidden states in the normal way.'     )
   def 	__UpperCAmelCase   (      self      :			Tuple     )  ->			Union[str, Any]:
        pass
   @unittest.skip('ESMfold does not output hidden states in the normal way.'     )
   def 	__UpperCAmelCase   (      self      :			Tuple     )  ->			Optional[Any]:
        pass
   @unittest.skip('ESMFold only has one output format.'     )
   def 	__UpperCAmelCase   (      self      :			Tuple     )  ->			Any:
        pass
   @unittest.skip('This test doesn\'t work for ESMFold and doesn\'t test core functionality'     )
   def 	__UpperCAmelCase   (      self      :			List[str]     )  ->			Union[str, Any]:
        pass
   @unittest.skip('ESMFold does not support input chunking.'     )
   def 	__UpperCAmelCase   (      self      :			int     )  ->			int:
        pass
   @unittest.skip('ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.'     )
   def 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			str:
        pass
   @unittest.skip('ESMFold doesn\'t support torchscript compilation.'     )
   def 	__UpperCAmelCase   (      self      :			str     )  ->			str:
        pass
   @unittest.skip('ESMFold doesn\'t support torchscript compilation.'     )
   def 	__UpperCAmelCase   (      self      :			str     )  ->			Tuple:
        pass
   @unittest.skip('ESMFold doesn\'t support torchscript compilation.'     )
   def 	__UpperCAmelCase   (      self      :			List[Any]     )  ->			Optional[Any]:
        pass
   @unittest.skip('ESMFold doesn\'t support data parallel.'     )
   def 	__UpperCAmelCase   (      self      :			Tuple     )  ->			Optional[int]:
        pass

   @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'     )
   def 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			List[Any]:
        pass







@require_torch
class lowercase_       (		_UpperCamelCase			):







   """simple docstring"""

   @slow
   def 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			str:
        _A    		=  EsmForProteinFolding.from_pretrained('facebook/esmfold_v1'     ).float()
        model.eval()
        _A    		=  torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]     )
        _A    		=  model(UpperCamelCase__     )['positions']
        _A    		=  torch.tensor([2.5_828, 0.7_993, -10.9_334], dtype=torch.floataa     )
        self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0], UpperCamelCase__, atol=1e-4     )     )








"""simple docstring"""







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  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Optional[int]     ,		lowerCamelCase_:							Union[str, Any] ):
						"""simple docstring"""






						if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 3:
												# expert layer
												snake_case   :			List[str]						    = flax_key_tuple[:-1] + ("weight",)
												snake_case   :			Union[str, Any]						    = torch.permute(lowerCamelCase_     ,		(0, 2, 1) )
						elif flax_key_tuple[-1] == "kernel" and ".".join(lowerCamelCase_ ):
												# linear layer
												snake_case   :			int						    = flax_key_tuple[:-1] + ("weight",)
												snake_case   :			Dict						    = flax_tensor.T
						elif flax_key_tuple[-1] in ["scale", "embedding"]:
												snake_case   :			Any						    = flax_key_tuple[:-1] + ("weight",)

						return flax_key_tuple, flax_tensor

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Union[str, Any]     ,		lowerCamelCase_:							Union[str, Any]     ,		lowerCamelCase_:							str ):
						"""simple docstring"""






						if "metadata" in layer:
												snake_case   :			Dict						    = layer.split("metadata" )
												snake_case   :			Optional[Any]						    = "".join(split_layer[0] )[:-1]
												snake_case   :			Any						    = [tuple(("metadata" + split_layer[1]).split("/" ) )]
						elif "kvstore" in layer:
												snake_case   :			List[str]						    = layer.split("kvstore" )
												snake_case   :			Tuple						    = "".join(split_layer[0] )[:-1]
												snake_case   :			Union[str, Any]						    = [tuple(("kvstore" + split_layer[1]).split("/" ) )]

						else:
												snake_case   :			List[Any]						    = layer.split("/" )
												snake_case   :			Union[str, Any]						    = "/".join(split_layer[:-1] )
												snake_case   :			int						    = (split_layer[-1],)

						if "kvstore/path" in layer:
												snake_case   :			str						    = f'''{switch_checkpoint_path}/{checkpoint_info[layer]}'''
						elif "kvstore/driver" in layer:
												snake_case   :			Tuple						    = "file"
						else:
												snake_case   :			int						    = checkpoint_info[layer]

						return curr_real_layer_name, split_layer, content

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Optional[Any]     ,		lowerCamelCase_:							Dict ):
						"""simple docstring"""






						snake_case   :			Optional[int]						    = rename_keys(lowerCamelCase_ )
						snake_case   :			str						    = {}
						for k, v in current_block.items():
												snake_case   :			List[str]						    = v
						snake_case   :			List[str]						    = new_current_block
						torch.save(lowerCamelCase_     ,		lowerCamelCase_ )

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Tuple     ,		lowerCamelCase_:							Optional[Any]     ,		lowerCamelCase_:							Dict     ,		lowerCamelCase_:							int     ,		lowerCamelCase_:							str = WEIGHTS_NAME ):
						"""simple docstring"""






						snake_case   :			List[str]						    = convert_file_size_to_int(lowerCamelCase_ )
						snake_case   :			List[Any]						    = []
						snake_case   :			Dict						    = {}
						snake_case   :			str						    = 0
						snake_case   :			List[str]						    = 0

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

						snake_case   :			Optional[int]						    = {}
						for layer in checkpoint_info.keys():
												snake_case    ,			snake_case    ,			snake_case   :			Union[str, Any]						    = get_key_and_tensorstore_dict(
												    lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_ )
												if curr_real_layer_name in all_layers:
																		snake_case   :			str						    = content
												else:
																		snake_case   :			Any						    = {split_layer[-1]: content}

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

												# use the renaming pattern from the small conversion scripts
												snake_case    ,			snake_case   :			Dict						    = rename_base_flax_keys(tuple(key.split("/" ) )     ,		lowerCamelCase_ )
												snake_case   :			Union[str, Any]						    = "/".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:
																		snake_case   :			str						    = 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
																		snake_case   :			Any						    = {}
																		snake_case   :			Union[str, Any]						    = 0

												snake_case   :			Any						    = raw_weights.to(getattr(lowerCamelCase_     ,		lowerCamelCase_ ) )
												current_block_size += weight_size
												total_size += weight_size

						# Add the last block
						snake_case   :			List[Any]						    = 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
						snake_case   :			List[Any]						    = {}
						snake_case   :			Dict						    = {}
						for idx, shard in enumerate(lowerCamelCase_ ):
												snake_case   :			List[Any]						    = weights_name.replace(
												    ".bin"     ,		f'''-{idx+1:05d}-of-{len(lowerCamelCase_ ):05d}.bin''' )  # len(sharded_state_dicts):05d}
												snake_case   :			Tuple						    = os.path.join(lowerCamelCase_     ,		weights_name.replace(".bin"     ,		f'''-{idx+1:05d}-of-???.bin''' ) )
												os.rename(lowerCamelCase_     ,		os.path.join(lowerCamelCase_     ,		lowerCamelCase_ ) )
												snake_case   :			Union[str, Any]						    = shard
												for key in shard:
																		snake_case   :			List[Any]						    = shard_file

    # Add the metadata
						snake_case   :			Optional[int]						    = {"total_size": total_size}
						snake_case   :			List[str]						    = {"metadata": metadata, "weight_map": weight_map}

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

						return metadata, index


if __name__ == "__main__":
			A		=						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.',
			)
			A		=						parser.parse_args()
			shard_on_the_fly(
			    args.switch_tax_checkpoint_path,
			    args.pytorch_dump_folder_path,
			    args.max_shard_size,
			    args.dtype,
			)

def  __SCREAMING_SNAKE_CASE		(			):
						"""simple docstring"""






						from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration, TaTokenizer

						snake_case   :			List[Any]						    = SwitchTransformersConfig.from_pretrained("google/switch-base-8" )
						config.save_pretrained("/home/arthur_huggingface_co/transformers/switch_converted" )
						snake_case   :			List[Any]						    = SwitchTransformersForConditionalGeneration.from_pretrained(
						    "/home/arthur_huggingface_co/transformers/switch_converted"     ,		device_map="auto" )

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

						snake_case   :			Dict						    = tokenizer(lowerCamelCase_     ,		return_tensors="pt" ).input_ids
						snake_case   :			Optional[int]						    = model.generate(lowerCamelCase_     ,		decoder_start_token_id=0 )
						print(tokenizer.decode(out[0] ) )



