import argparse
import json

import torch

from diffusers import DDPMScheduler, LDMPipeline, UNetaDModel, VQModel



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE=1			):
    '''simple docstring'''







    if n_shave_prefix_segments >= 0:
        return ".".join(path.split('''.'''			)[n_shave_prefix_segments:]			)
    else:
        return ".".join(path.split('''.'''			)[:n_shave_prefix_segments]			)



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE=0			):
    '''simple docstring'''







    __UpperCamelCase  :Optional[Any]    						=	[]
    for old_item in old_list:
        __UpperCamelCase  :Union[str, Any]    						=	old_item.replace('''in_layers.0'''						,    '''norm1'''			)
        __UpperCamelCase  :Dict    						=	new_item.replace('''in_layers.2'''						,    '''conv1'''			)

        __UpperCamelCase  :Optional[Any]    						=	new_item.replace('''out_layers.0'''						,    '''norm2'''			)
        __UpperCamelCase  :int    						=	new_item.replace('''out_layers.3'''						,    '''conv2'''			)

        __UpperCamelCase  :Dict    						=	new_item.replace('''emb_layers.1'''						,    '''time_emb_proj'''			)
        __UpperCamelCase  :Any    						=	new_item.replace('''skip_connection'''						,    '''conv_shortcut'''			)

        __UpperCamelCase  :List[str]    						=	shave_segments(SCREAMING_SNAKE_CASE						,    n_shave_prefix_segments=SCREAMING_SNAKE_CASE			)

        mapping.append({'''old''': old_item, '''new''': new_item}			)

    return mapping



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE=0			):
    '''simple docstring'''







    __UpperCamelCase  :Optional[Any]    						=	[]
    for old_item in old_list:
        __UpperCamelCase  :List[Any]    						=	old_item

        __UpperCamelCase  :int    						=	new_item.replace('''norm.weight'''						,    '''group_norm.weight'''			)
        __UpperCamelCase  :Optional[int]    						=	new_item.replace('''norm.bias'''						,    '''group_norm.bias'''			)

        __UpperCamelCase  :List[str]    						=	new_item.replace('''proj_out.weight'''						,    '''proj_attn.weight'''			)
        __UpperCamelCase  :str    						=	new_item.replace('''proj_out.bias'''						,    '''proj_attn.bias'''			)

        __UpperCamelCase  :Dict    						=	shave_segments(SCREAMING_SNAKE_CASE						,    n_shave_prefix_segments=SCREAMING_SNAKE_CASE			)

        mapping.append({'''old''': old_item, '''new''': new_item}			)

    return mapping



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE=None						,    SCREAMING_SNAKE_CASE=None						,    SCREAMING_SNAKE_CASE=None			):
    '''simple docstring'''







    assert isinstance(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			), "Paths should be a list of dicts containing 'old' and 'new' keys."

    # Splits the attention layers into three variables.
    if attention_paths_to_split is not None:
        for path, path_map in attention_paths_to_split.items():
            __UpperCamelCase  :Tuple    						=	old_checkpoint[path]
            __UpperCamelCase  :Optional[Any]    						=	old_tensor.shape[0] // 3

            __UpperCamelCase  :Optional[int]    						=	(-1, channels) if len(old_tensor.shape			) == 3 else (-1)

            __UpperCamelCase  :Any    						=	old_tensor.shape[0] // config['''num_head_channels'''] // 3

            __UpperCamelCase  :Dict    						=	old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]			)
            __UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase  :int    						=	old_tensor.split(channels // num_heads						,    dim=1			)

            __UpperCamelCase  :Union[str, Any]    						=	query.reshape(SCREAMING_SNAKE_CASE			)
            __UpperCamelCase  :Dict    						=	key.reshape(SCREAMING_SNAKE_CASE			)
            __UpperCamelCase  :int    						=	value.reshape(SCREAMING_SNAKE_CASE			)

    for path in paths:
        __UpperCamelCase  :int    						=	path['''new''']

        # These have already been assigned
        if attention_paths_to_split is not None and new_path in attention_paths_to_split:
            continue

        # Global renaming happens here
        __UpperCamelCase  :Any    						=	new_path.replace('''middle_block.0'''						,    '''mid_block.resnets.0'''			)
        __UpperCamelCase  :List[Any]    						=	new_path.replace('''middle_block.1'''						,    '''mid_block.attentions.0'''			)
        __UpperCamelCase  :List[Any]    						=	new_path.replace('''middle_block.2'''						,    '''mid_block.resnets.1'''			)

        if additional_replacements is not None:
            for replacement in additional_replacements:
                __UpperCamelCase  :Optional[Any]    						=	new_path.replace(replacement['''old''']						,    replacement['''new''']			)

        # proj_attn.weight has to be converted from conv 1D to linear
        if "proj_attn.weight" in new_path:
            __UpperCamelCase  :str    						=	old_checkpoint[path['''old''']][:, :, 0]
        else:
            __UpperCamelCase  :int    						=	old_checkpoint[path['''old''']]



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :List[str]    						=	{}

    __UpperCamelCase  :Optional[int]    						=	checkpoint['''time_embed.0.weight''']
    __UpperCamelCase  :Union[str, Any]    						=	checkpoint['''time_embed.0.bias''']
    __UpperCamelCase  :str    						=	checkpoint['''time_embed.2.weight''']
    __UpperCamelCase  :List[str]    						=	checkpoint['''time_embed.2.bias''']

    __UpperCamelCase  :Optional[Any]    						=	checkpoint['''input_blocks.0.0.weight''']
    __UpperCamelCase  :Dict    						=	checkpoint['''input_blocks.0.0.bias''']

    __UpperCamelCase  :int    						=	checkpoint['''out.0.weight''']
    __UpperCamelCase  :str    						=	checkpoint['''out.0.bias''']
    __UpperCamelCase  :Optional[int]    						=	checkpoint['''out.2.weight''']
    __UpperCamelCase  :Any    						=	checkpoint['''out.2.bias''']

    # Retrieves the keys for the input blocks only
    __UpperCamelCase  :List[str]    						=	len({'''.'''.join(layer.split('''.'''			)[:2]			) for layer in checkpoint if '''input_blocks''' in layer}			)
    __UpperCamelCase  :str    						=	{
        layer_id: [key for key in checkpoint if f"""input_blocks.{layer_id}""" in key]
        for layer_id in range(SCREAMING_SNAKE_CASE			)
    }

    # Retrieves the keys for the middle blocks only
    __UpperCamelCase  :Union[str, Any]    						=	len({'''.'''.join(layer.split('''.'''			)[:2]			) for layer in checkpoint if '''middle_block''' in layer}			)
    __UpperCamelCase  :List[Any]    						=	{
        layer_id: [key for key in checkpoint if f"""middle_block.{layer_id}""" in key]
        for layer_id in range(SCREAMING_SNAKE_CASE			)
    }

    # Retrieves the keys for the output blocks only
    __UpperCamelCase  :Tuple    						=	len({'''.'''.join(layer.split('''.'''			)[:2]			) for layer in checkpoint if '''output_blocks''' in layer}			)
    __UpperCamelCase  :List[str]    						=	{
        layer_id: [key for key in checkpoint if f"""output_blocks.{layer_id}""" in key]
        for layer_id in range(SCREAMING_SNAKE_CASE			)
    }

    for i in range(1						,    SCREAMING_SNAKE_CASE			):
        __UpperCamelCase  :Any    						=	(i - 1) // (config['''num_res_blocks'''] + 1)
        __UpperCamelCase  :Tuple    						=	(i - 1) % (config['''num_res_blocks'''] + 1)

        __UpperCamelCase  :Dict    						=	[key for key in input_blocks[i] if f"""input_blocks.{i}.0""" in key]
        __UpperCamelCase  :Tuple    						=	[key for key in input_blocks[i] if f"""input_blocks.{i}.1""" in key]

        if f"""input_blocks.{i}.0.op.weight""" in checkpoint:
            __UpperCamelCase  :Optional[Any]    						=	checkpoint[
                f"""input_blocks.{i}.0.op.weight"""
            ]
            __UpperCamelCase  :Tuple    						=	checkpoint[
                f"""input_blocks.{i}.0.op.bias"""
            ]
            continue

        __UpperCamelCase  :List[Any]    						=	renew_resnet_paths(SCREAMING_SNAKE_CASE			)
        __UpperCamelCase  :Union[str, Any]    						=	{'''old''': f"""input_blocks.{i}.0""", '''new''': f"""down_blocks.{block_id}.resnets.{layer_in_block_id}"""}
        __UpperCamelCase  :int    						=	{'''old''': '''resnets.2.op''', '''new''': '''downsamplers.0.op'''}
        assign_to_checkpoint(
            SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    additional_replacements=[meta_path, resnet_op]						,    config=SCREAMING_SNAKE_CASE			)

        if len(SCREAMING_SNAKE_CASE			):
            __UpperCamelCase  :Dict    						=	renew_attention_paths(SCREAMING_SNAKE_CASE			)
            __UpperCamelCase  :Any    						=	{
                '''old''': f"""input_blocks.{i}.1""",
                '''new''': f"""down_blocks.{block_id}.attentions.{layer_in_block_id}""",
            }
            __UpperCamelCase  :Dict    						=	{
                f"""input_blocks.{i}.1.qkv.bias""": {
                    '''key''': f"""down_blocks.{block_id}.attentions.{layer_in_block_id}.key.bias""",
                    '''query''': f"""down_blocks.{block_id}.attentions.{layer_in_block_id}.query.bias""",
                    '''value''': f"""down_blocks.{block_id}.attentions.{layer_in_block_id}.value.bias""",
                },
                f"""input_blocks.{i}.1.qkv.weight""": {
                    '''key''': f"""down_blocks.{block_id}.attentions.{layer_in_block_id}.key.weight""",
                    '''query''': f"""down_blocks.{block_id}.attentions.{layer_in_block_id}.query.weight""",
                    '''value''': f"""down_blocks.{block_id}.attentions.{layer_in_block_id}.value.weight""",
                },
            }
            assign_to_checkpoint(
                SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    additional_replacements=[meta_path]						,    attention_paths_to_split=SCREAMING_SNAKE_CASE						,    config=SCREAMING_SNAKE_CASE						,    )

    __UpperCamelCase  :Union[str, Any]    						=	middle_blocks[0]
    __UpperCamelCase  :str    						=	middle_blocks[1]
    __UpperCamelCase  :Dict    						=	middle_blocks[2]

    __UpperCamelCase  :List[str]    						=	renew_resnet_paths(SCREAMING_SNAKE_CASE			)
    assign_to_checkpoint(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    config=SCREAMING_SNAKE_CASE			)

    __UpperCamelCase  :Tuple    						=	renew_resnet_paths(SCREAMING_SNAKE_CASE			)
    assign_to_checkpoint(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    config=SCREAMING_SNAKE_CASE			)

    __UpperCamelCase  :Optional[int]    						=	renew_attention_paths(SCREAMING_SNAKE_CASE			)
    __UpperCamelCase  :Union[str, Any]    						=	{
        '''middle_block.1.qkv.bias''': {
            '''key''': '''mid_block.attentions.0.key.bias''',
            '''query''': '''mid_block.attentions.0.query.bias''',
            '''value''': '''mid_block.attentions.0.value.bias''',
        },
        '''middle_block.1.qkv.weight''': {
            '''key''': '''mid_block.attentions.0.key.weight''',
            '''query''': '''mid_block.attentions.0.query.weight''',
            '''value''': '''mid_block.attentions.0.value.weight''',
        },
    }
    assign_to_checkpoint(
        SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    attention_paths_to_split=SCREAMING_SNAKE_CASE						,    config=SCREAMING_SNAKE_CASE			)

    for i in range(SCREAMING_SNAKE_CASE			):
        __UpperCamelCase  :Optional[Any]    						=	i // (config['''num_res_blocks'''] + 1)
        __UpperCamelCase  :Tuple    						=	i % (config['''num_res_blocks'''] + 1)
        __UpperCamelCase  :Optional[Any]    						=	[shave_segments(SCREAMING_SNAKE_CASE						,    2			) for name in output_blocks[i]]
        __UpperCamelCase  :str    						=	{}

        for layer in output_block_layers:
            __UpperCamelCase						, __UpperCamelCase  :Optional[int]    						=	layer.split('''.'''			)[0], shave_segments(SCREAMING_SNAKE_CASE						,    1			)
            if layer_id in output_block_list:
                output_block_list[layer_id].append(SCREAMING_SNAKE_CASE			)
            else:
                __UpperCamelCase  :Union[str, Any]    						=	[layer_name]

        if len(SCREAMING_SNAKE_CASE			) > 1:
            __UpperCamelCase  :Union[str, Any]    						=	[key for key in output_blocks[i] if f"""output_blocks.{i}.0""" in key]
            __UpperCamelCase  :int    						=	[key for key in output_blocks[i] if f"""output_blocks.{i}.1""" in key]

            __UpperCamelCase  :List[Any]    						=	renew_resnet_paths(SCREAMING_SNAKE_CASE			)
            __UpperCamelCase  :Dict    						=	renew_resnet_paths(SCREAMING_SNAKE_CASE			)

            __UpperCamelCase  :Any    						=	{'''old''': f"""output_blocks.{i}.0""", '''new''': f"""up_blocks.{block_id}.resnets.{layer_in_block_id}"""}
            assign_to_checkpoint(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    additional_replacements=[meta_path]						,    config=SCREAMING_SNAKE_CASE			)

            if ["conv.weight", "conv.bias"] in output_block_list.values():
                __UpperCamelCase  :int    						=	list(output_block_list.values()			).index(['''conv.weight''', '''conv.bias''']			)
                __UpperCamelCase  :List[Any]    						=	checkpoint[
                    f"""output_blocks.{i}.{index}.conv.weight"""
                ]
                __UpperCamelCase  :List[str]    						=	checkpoint[
                    f"""output_blocks.{i}.{index}.conv.bias"""
                ]

                # Clear attentions as they have been attributed above.
                if len(SCREAMING_SNAKE_CASE			) == 2:
                    __UpperCamelCase  :Tuple    						=	[]

            if len(SCREAMING_SNAKE_CASE			):
                __UpperCamelCase  :Tuple    						=	renew_attention_paths(SCREAMING_SNAKE_CASE			)
                __UpperCamelCase  :Tuple    						=	{
                    '''old''': f"""output_blocks.{i}.1""",
                    '''new''': f"""up_blocks.{block_id}.attentions.{layer_in_block_id}""",
                }
                __UpperCamelCase  :Optional[Any]    						=	{
                    f"""output_blocks.{i}.1.qkv.bias""": {
                        '''key''': f"""up_blocks.{block_id}.attentions.{layer_in_block_id}.key.bias""",
                        '''query''': f"""up_blocks.{block_id}.attentions.{layer_in_block_id}.query.bias""",
                        '''value''': f"""up_blocks.{block_id}.attentions.{layer_in_block_id}.value.bias""",
                    },
                    f"""output_blocks.{i}.1.qkv.weight""": {
                        '''key''': f"""up_blocks.{block_id}.attentions.{layer_in_block_id}.key.weight""",
                        '''query''': f"""up_blocks.{block_id}.attentions.{layer_in_block_id}.query.weight""",
                        '''value''': f"""up_blocks.{block_id}.attentions.{layer_in_block_id}.value.weight""",
                    },
                }
                assign_to_checkpoint(
                    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    additional_replacements=[meta_path]						,    attention_paths_to_split=to_split if any('''qkv''' in key for key in attentions			) else None						,    config=SCREAMING_SNAKE_CASE						,    )
        else:
            __UpperCamelCase  :Dict    						=	renew_resnet_paths(SCREAMING_SNAKE_CASE						,    n_shave_prefix_segments=1			)
            for path in resnet_0_paths:
                __UpperCamelCase  :Union[str, Any]    						=	'''.'''.join(['''output_blocks''', str(SCREAMING_SNAKE_CASE			), path['''old''']]			)
                __UpperCamelCase  :str    						=	'''.'''.join(['''up_blocks''', str(SCREAMING_SNAKE_CASE			), '''resnets''', str(SCREAMING_SNAKE_CASE			), path['''new''']]			)

                __UpperCamelCase  :List[Any]    						=	checkpoint[old_path]

    return new_checkpoint


if __name__ == "__main__":
    __lowercase											=      argparse.ArgumentParser()

    parser.add_argument(
        '''--checkpoint_path''', default=None, type=str, required=True, help='''Path to the checkpoint to convert.'''
    )

    parser.add_argument(
        '''--config_file''',
        default=None,
        type=str,
        required=True,
        help='''The config json file corresponding to the architecture.''',
    )

    parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the output model.''')

    __lowercase											=      parser.parse_args()

    __lowercase											=      torch.load(args.checkpoint_path)

    with open(args.config_file) as f:
        __lowercase											=      json.loads(f.read())

    __lowercase											=      convert_ldm_checkpoint(checkpoint, config)

    if "ldm" in config:
        del config["ldm"]

    __lowercase											=      UNetaDModel(**config)
    model.load_state_dict(converted_checkpoint)

    try:
        __lowercase											=      DDPMScheduler.from_config('''/'''.join(args.checkpoint_path.split('''/''')[:-1]))
        __lowercase											=      VQModel.from_pretrained('''/'''.join(args.checkpoint_path.split('''/''')[:-1]))

        __lowercase											=      LDMPipeline(unet=model, scheduler=scheduler, vae=vqvae)
        pipe.save_pretrained(args.dump_path)
    except:  # noqa: E722
        model.save_pretrained(args.dump_path)






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 			lowerCamelCase_  :






    '''simple docstring'''






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







    def     UpperCamelCase__     (			self) ->			List[Any]:
        __UpperCamelCase  :int    						=	ids_tensor([self.batch_size, self.seq_length]				,				self.vocab_size)

        __UpperCamelCase  :int    						=	None
        if self.use_input_mask:
            __UpperCamelCase  :Optional[Any]    						=	random_attention_mask([self.batch_size, self.seq_length])

        __UpperCamelCase  :Dict    						=	None
        __UpperCamelCase  :List[Any]    						=	None
        __UpperCamelCase  :Tuple    						=	None
        if self.use_labels:
            __UpperCamelCase  :Union[str, Any]    						=	ids_tensor([self.batch_size]				,				self.type_sequence_label_size)
            __UpperCamelCase  :Union[str, Any]    						=	ids_tensor([self.batch_size, self.seq_length]				,				self.num_labels)
            __UpperCamelCase  :str    						=	ids_tensor([self.batch_size]				,				self.num_choices)

        __UpperCamelCase  :Optional[Any]    						=	self.get_config()

        return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels







    def     UpperCamelCase__     (			self) ->			Optional[int]:
        __UpperCamelCase  :int    						=	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=__lowercase				,				esmfold_config={'''trunk''': {'''num_blocks''': 2}, '''fp16_esm''': False}				,				)
        return config







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase) ->			Union[str, Any]:
        __UpperCamelCase  :int    						=	EsmForProteinFolding(config=__lowercase).float()
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase  :Tuple    						=	model(__lowercase				,				attention_mask=__lowercase)
        __UpperCamelCase  :Any    						=	model(__lowercase)
        __UpperCamelCase  :Optional[Any]    						=	model(__lowercase)

        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) ->			Optional[Any]:
        __UpperCamelCase  :Dict    						=	self.prepare_config_and_inputs()
        (
            (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, 
        )  :List[str]    						=	config_and_inputs
        __UpperCamelCase  :Tuple    						=	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
        return config, inputs_dict


@require_torch
class 			lowerCamelCase_  (     UpperCAmelCase_		,			UpperCAmelCase_		,			unittest.TestCase				):






    '''simple docstring'''






    a__			:		Optional[int]      			=   False

    a__			:		Optional[int]      			=   (EsmForProteinFolding,) if is_torch_available() else ()
    a__			:		str      			=   ()
    a__			:		Tuple      			=   {} if is_torch_available() else {}
    a__			:		List[Any]      			=   False







    def     UpperCamelCase__     (			self) ->			Tuple:
        __UpperCamelCase  :Dict    						=	EsmFoldModelTester(self)
        __UpperCamelCase  :Dict    						=	ConfigTester(self				,				config_class=__lowercase				,				hidden_size=37)







    def     UpperCamelCase__     (			self) ->			List[Any]:
        self.config_tester.run_common_tests()







    def     UpperCamelCase__     (			self) ->			List[str]:
        __UpperCamelCase  :Any    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowercase)







    @unittest.skip('''Does not support attention outputs''')
    def     UpperCamelCase__     (			self) ->			Any:
        pass







    @unittest.skip
    def     UpperCamelCase__     (			self) ->			Any:
        pass







    @unittest.skip('''Esm does not support embedding resizing''')
    def     UpperCamelCase__     (			self) ->			Union[str, Any]:
        pass







    @unittest.skip('''Esm does not support embedding resizing''')
    def     UpperCamelCase__     (			self) ->			Optional[Any]:
        pass







    @unittest.skip('''ESMFold does not support passing input embeds!''')
    def     UpperCamelCase__     (			self) ->			List[Any]:
        pass







    @unittest.skip('''ESMFold does not support head pruning.''')
    def     UpperCamelCase__     (			self) ->			Optional[int]:
        pass







    @unittest.skip('''ESMFold does not support head pruning.''')
    def     UpperCamelCase__     (			self) ->			Optional[Any]:
        pass







    @unittest.skip('''ESMFold does not support head pruning.''')
    def     UpperCamelCase__     (			self) ->			Dict:
        pass







    @unittest.skip('''ESMFold does not support head pruning.''')
    def     UpperCamelCase__     (			self) ->			Optional[Any]:
        pass







    @unittest.skip('''ESMFold does not support head pruning.''')
    def     UpperCamelCase__     (			self) ->			Optional[int]:
        pass







    @unittest.skip('''ESMFold does not output hidden states in the normal way.''')
    def     UpperCamelCase__     (			self) ->			str:
        pass







    @unittest.skip('''ESMfold does not output hidden states in the normal way.''')
    def     UpperCamelCase__     (			self) ->			Optional[Any]:
        pass







    @unittest.skip('''ESMFold only has one output format.''')
    def     UpperCamelCase__     (			self) ->			Union[str, Any]:
        pass







    @unittest.skip('''This test doesn\'t work for ESMFold and doesn\'t test core functionality''')
    def     UpperCamelCase__     (			self) ->			List[Any]:
        pass







    @unittest.skip('''ESMFold does not support input chunking.''')
    def     UpperCamelCase__     (			self) ->			List[Any]:
        pass







    @unittest.skip('''ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.''')
    def     UpperCamelCase__     (			self) ->			int:
        pass







    @unittest.skip('''ESMFold doesn\'t support torchscript compilation.''')
    def     UpperCamelCase__     (			self) ->			str:
        pass







    @unittest.skip('''ESMFold doesn\'t support torchscript compilation.''')
    def     UpperCamelCase__     (			self) ->			List[Any]:
        pass







    @unittest.skip('''ESMFold doesn\'t support torchscript compilation.''')
    def     UpperCamelCase__     (			self) ->			List[Any]:
        pass







    @unittest.skip('''ESMFold doesn\'t support data parallel.''')
    def     UpperCamelCase__     (			self) ->			Any:
        pass







    @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.''')
    def     UpperCamelCase__     (			self) ->			Dict:
        pass


@require_torch
class 			lowerCamelCase_  (     UpperCAmelCase_				):






    '''simple docstring'''






    @slow
    def     UpperCamelCase__     (			self) ->			Tuple:
        __UpperCamelCase  :Optional[Any]    						=	EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1''').float()
        model.eval()
        __UpperCamelCase  :Tuple    						=	torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]])
        __UpperCamelCase  :List[Any]    						=	model(__lowercase)['''positions''']
        __UpperCamelCase  :Optional[int]    						=	torch.tensor([2.58_28, 0.79_93, -10.93_34]				,				dtype=torch.floataa)
        self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]				,				__lowercase				,				atol=1E-4))





