






"""simple docstring"""







import argparse
import pathlib
from pathlib import Path
from tempfile import TemporaryDirectory

import esm as esm_module
import torch
from esm.esmfold.va.misc import batch_encode_sequences as esmfold_encode_sequences
from esm.esmfold.va.pretrained import esmfold_va

from transformers.models.esm.configuration_esm import EsmConfig, EsmFoldConfig
from transformers.models.esm.modeling_esm import (
    EsmForMaskedLM,
    EsmForSequenceClassification,
    EsmIntermediate,
    EsmLayer,
    EsmOutput,
    EsmSelfAttention,
    EsmSelfOutput,
)
from transformers.models.esm.modeling_esmfold import EsmForProteinFolding
from transformers.models.esm.tokenization_esm import EsmTokenizer
from transformers.utils import logging


logging.set_verbosity_info()
__a							:						List[str]   							=			logging.get_logger(__name__)

__a							:						Any   							=			[
    (
        'protein1',
        'MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA',
    ),
    ('protein2', 'MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA'),
    ('protein3', 'MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG'),
    ('protein4', 'MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA'),
]

__a							:						Tuple   							=			{
    'esm1b_t33_650M_UR50S': esm_module.pretrained.esm1b_t33_650M_UR50S,
    'esm1v_t33_650M_UR90S_1': esm_module.pretrained.esm1v_t33_650M_UR90S_1,
    'esm1v_t33_650M_UR90S_2': esm_module.pretrained.esm1v_t33_650M_UR90S_2,
    'esm1v_t33_650M_UR90S_3': esm_module.pretrained.esm1v_t33_650M_UR90S_3,
    'esm1v_t33_650M_UR90S_4': esm_module.pretrained.esm1v_t33_650M_UR90S_4,
    'esm1v_t33_650M_UR90S_5': esm_module.pretrained.esm1v_t33_650M_UR90S_5,
    'esm2_t48_15B_UR50D': esm_module.pretrained.esm2_t48_15B_UR50D,
    'esm2_t36_3B_UR50D': esm_module.pretrained.esm2_t36_3B_UR50D,
    'esm2_t33_650M_UR50D': esm_module.pretrained.esm2_t33_650M_UR50D,
    'esm2_t30_150M_UR50D': esm_module.pretrained.esm2_t30_150M_UR50D,
    'esm2_t12_35M_UR50D': esm_module.pretrained.esm2_t12_35M_UR50D,
    'esm2_t6_8M_UR50D': esm_module.pretrained.esm2_t6_8M_UR50D,
    'esmfold_v1': esmfold_va,
}

__a							:						Optional[Any]   							=			list('ARNDCQEGHILKMFPSTWYV')

__a							:						List[Any]   							=			restypes + ['X']
__a							:						Optional[Any]   							=			restypes_with_x + ['<pad>', '<mask>', '<cls>', '<sep>', '<eos>']





def 							SCREAMING_SNAKE_CASE							(   ):
					with TemporaryDirectory() as tempdir:
										a__ 							=	'\n'.join(lowercase_)
										a__ 							=	Path(lowercase_) / 'vocab.txt'
										vocab_file.write_text(lowercase_)
										a__ 							=	EsmTokenizer(vocab_file=str(lowercase_))
					a__ 							=	0  # Overlaps with 'A' but that seems to be what they want
					return hf_tokenizer





def 							SCREAMING_SNAKE_CASE							(   lowerCamelCase_			,  lowerCamelCase_):
					a__ 							=	our_module.load_state_dict(original_module.state_dict())
					if status.missing_keys:
										raise ValueError(f'Missing keys: {status.missing_keys}')
					if status.unexpected_keys:
										raise ValueError(f'Unexpected keys: {status.unexpected_keys}')





def 							SCREAMING_SNAKE_CASE							(   lowerCamelCase_			,  lowerCamelCase_			,  lowerCamelCase_			,  lowerCamelCase_			,  lowerCamelCase_):
					if model.startswith('''esmfold'''):
										a__ 							=	MODEL_MAPPING[model]()
					else:
										a__ 							=	MODEL_MAPPING[model]()
					esm.eval()  # disable dropout

					if model.startswith('''esmfold'''):
										a__ 							=	esm.esm.embed_dim
										a__ 							=	esm.esm.num_layers
										a__ 							=	esm.esm.attention_heads
										a__ 							=	4 * embed_dim
										a__ 							=	esm.esm.token_dropout
										a__ 							=	False  # This code path does not exist in ESM-2
										a__ 							=	'rotary'
										a__ 							=	True
										a__ 							=	EsmFoldConfig()
										for key, val in esm.cfg.items():
															if hasattr(lowercase_			,  lowercase_) and key != "trunk":
																				setattr(lowercase_			,  lowercase_			,  lowercase_)
										for key, val in esm.cfg.trunk.items():
															if hasattr(esmfold_config.trunk			,  lowercase_) and key != "structure_module":
																				setattr(esmfold_config.trunk			,  lowercase_			,  lowercase_)
										for key, val in esm.cfg.trunk.structure_module.items():
															if hasattr(esmfold_config.trunk.structure_module			,  lowercase_):
																				setattr(esmfold_config.trunk.structure_module			,  lowercase_			,  lowercase_)
					elif hasattr(lowercase_			,  '''args'''):
										# Indicates an ESM-1b or ESM-1v model
										a__ 							=	esm.args.embed_dim
										a__ 							=	esm.args.layers
										a__ 							=	esm.args.attention_heads
										a__ 							=	esm.args.ffn_embed_dim
										a__ 							=	esm.args.token_dropout
										a__ 							=	True if esm.emb_layer_norm_before else False
										a__ 							=	'absolute'
										a__ 							=	False
										a__ 							=	None
					else:
										# Indicates an ESM-2 model
										a__ 							=	esm.embed_dim
										a__ 							=	esm.num_layers
										a__ 							=	esm.attention_heads
										a__ 							=	4 * embed_dim  # This is hardcoded in ESM-2
										a__ 							=	esm.token_dropout
										a__ 							=	False  # This code path does not exist in ESM-2
										a__ 							=	'rotary'
										a__ 							=	False
										a__ 							=	None

					if is_folding_model:
										a__ 							=	esm.esm.alphabet
					a__ 							=	tuple(alphabet.all_toks)
					a__ 							=	alphabet.mask_idx
					a__ 							=	alphabet.padding_idx

					if is_folding_model:
										a__ 							=	esm.esm
					else:
										a__ 							=	esm

					a__ 							=	EsmConfig(
					    vocab_size=original_esm_model.embed_tokens.num_embeddings			,  mask_token_id=lowercase_			,  hidden_size=lowercase_			,  num_hidden_layers=lowercase_			,  num_attention_heads=lowercase_			,  intermediate_size=lowercase_			,  max_position_embeddings=1026			,  layer_norm_eps=1E-5			,  attention_probs_dropout_prob=0.0			,  hidden_dropout_prob=0.0			,  pad_token_id=lowercase_			,  emb_layer_norm_before=lowercase_			,  token_dropout=lowercase_			,  position_embedding_type=lowercase_			,  is_folding_model=lowercase_			,  esmfold_config=lowercase_			,  vocab_list=lowercase_			,  )
					if classification_head:
										a__ 							=	esm.classification_heads['mnli'].out_proj.weight.shape[0]
					print('''Our ESM config:'''			,  lowercase_)

					if model.startswith('''esmfold'''):
										a__ 							=	EsmForProteinFolding
					elif classification_head:
										a__ 							=	EsmForSequenceClassification
					else:
										a__ 							=	EsmForMaskedLM
					a__ 							=	model_class(lowercase_)
					model.eval()

					# Now let's copy all the weights.
					# Embeddings
					a__ 							=	original_esm_model.embed_tokens.weight
					if position_embedding_type == "absolute":
										a__ 							=	original_esm_model.embed_positions.weight

					if config.emb_layer_norm_before:
										a__ 							=	original_esm_model.emb_layer_norm_before.weight
										a__ 							=	original_esm_model.emb_layer_norm_before.bias

					a__ 							=	original_esm_model.emb_layer_norm_after.weight
					a__ 							=	original_esm_model.emb_layer_norm_after.bias

					for i in range(config.num_hidden_layers):
										# Encoder: start of layer
										a__ 							=	model.esm.encoder.layer[i]
										# esm_layer: TransformerSentenceEncoderLayer = original_esm_model.layers[i]
										a__ 							=	original_esm_model.layers[i]

										# self attention
										a__ 							=	layer.attention.self
										assert (
										    esm_layer.self_attn.k_proj.weight.data.shape
										    == esm_layer.self_attn.q_proj.weight.data.shape
										    == esm_layer.self_attn.v_proj.weight.data.shape
										    == torch.Size((config.hidden_size, config.hidden_size))
										)

										a__ 							=	esm_layer.self_attn.q_proj.weight
										a__ 							=	esm_layer.self_attn.q_proj.bias
										a__ 							=	esm_layer.self_attn.k_proj.weight
										a__ 							=	esm_layer.self_attn.k_proj.bias
										a__ 							=	esm_layer.self_attn.v_proj.weight
										a__ 							=	esm_layer.self_attn.v_proj.bias

										if getattr(esm_layer.self_attn			,  '''rot_emb'''			,  lowercase_) is not None:
															# Matt: Although inv_freq is not a trainable weight, it is computed at model init and cached.
															# During the training of ESM-2 the model was converted to float16 precision, which also converts
															# the inv_freq tensor, and the loss of precision remains even if the model is loaded later as float32.
															# If we recompute inv_freq without this loss of precision then we will get subtly different rotary
															# embeddings, which are enough to cause significant discrepancies in model outputs. To avoid this,
															# we make sure the new model copies the data from the old inv_freq.
															a__ 							=	esm_layer.self_attn.rot_emb.inv_freq

										# LayerNorm changes for pre-activation
										a__ 							=	esm_layer.self_attn_layer_norm.weight
										a__ 							=	esm_layer.self_attn_layer_norm.bias
										a__ 							=	esm_layer.final_layer_norm.weight
										a__ 							=	esm_layer.final_layer_norm.bias

										# self-attention output
										a__ 							=	layer.attention.output
										assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
										a__ 							=	esm_layer.self_attn.out_proj.weight
										a__ 							=	esm_layer.self_attn.out_proj.bias

										# intermediate
										a__ 							=	layer.intermediate
										assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
										a__ 							=	esm_layer.fca.weight
										a__ 							=	esm_layer.fca.bias

										# output
										a__ 							=	layer.output
										assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
										a__ 							=	esm_layer.fca.weight
										a__ 							=	esm_layer.fca.bias
										# end of layer

					if is_folding_model:
										a__ 							=	esm.esm_s_combine.data
										a__ 							=	esm.afa_to_esm.data
										transfer_and_check_weights(esm.embedding			,  model.embedding)
										transfer_and_check_weights(esm.esm_s_mlp			,  model.esm_s_mlp)
										transfer_and_check_weights(esm.trunk			,  model.trunk)
										transfer_and_check_weights(esm.distogram_head			,  model.distogram_head)
										transfer_and_check_weights(esm.ptm_head			,  model.ptm_head)
										transfer_and_check_weights(esm.lm_head			,  model.lm_head)
										transfer_and_check_weights(esm.lddt_head			,  model.lddt_head)

					elif classification_head:
										a__ 							=	esm.esm.classification_heads['mnli'].dense.weight
										a__ 							=	esm.classification_heads['mnli'].dense.bias
										a__ 							=	esm.classification_heads['mnli'].out_proj.weight
										a__ 							=	esm.classification_heads['mnli'].out_proj.bias
					else:
										# LM Head
										a__ 							=	esm.lm_head.dense.weight
										a__ 							=	esm.lm_head.dense.bias
										a__ 							=	esm.lm_head.layer_norm.weight
										a__ 							=	esm.lm_head.layer_norm.bias
										a__ 							=	esm.lm_head.weight
										a__ 							=	esm.lm_head.bias

					# Contact prediction head
					transfer_and_check_weights(esm.contact_head			,  model.esm.contact_head)

					# Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
					if is_folding_model:
										# Folding models aren't trained on masked inputs and don't like mask tokens.
										a__ 							=	SAMPLE_DATA[:2]
					else:
										a__ 							=	SAMPLE_DATA

					if is_folding_model:
										a__ 							=	get_esmfold_tokenizer()
										a__ 							=	hf_tokenizer(
										    [row[1] for row in sample_data]			,  return_tensors='''pt'''			,  padding=lowercase_			,  add_special_tokens=lowercase_)
										a__ 							=	esmfold_encode_sequences([row[1] for row in sample_data])
										a__ 							=	torch.all(hf_tokens['''input_ids'''] == esmfold_aas) and torch.all(
										    hf_tokens['''attention_mask'''] == esmfold_mask)
					else:
										# Let's check that we get the same results.
										a__ 							=	alphabet.get_batch_converter()
										a__ 							=	batch_converter(lowercase_)
										# Prepare tokenizer and make sure it matches
										with TemporaryDirectory() as tempdir:
															a__ 							=	'\n'.join(alphabet.all_toks)
															a__ 							=	Path(lowercase_) / 'vocab.txt'
															vocab_file.write_text(lowercase_)
															a__ 							=	EsmTokenizer(vocab_file=str(lowercase_))

										a__ 							=	hf_tokenizer([row[1] for row in sample_data]			,  return_tensors='''pt'''			,  padding=lowercase_)
										a__ 							=	torch.all(hf_tokens['''input_ids'''] == batch_tokens)

					print('''Do both models tokenizers output the same tokens?'''			,  '''🔥''' if success else '''💩''')
					if not success:
										raise Exception('''Tokenization does not match!''')

					with torch.no_grad():
										if is_folding_model:
															# Let's test the model in parts
															# ESMFold always converts the ESM stem to float16, which requires float16 ops
															# that don't exist on CPU. Therefore, to test it we need to run it on GPU. However,
															# ESMFold is what we in the community call a "big boy" and so we desperately avoid putting both the
															# original and the converted model on the GPU at the same time.
															a__ 							=	esm.cuda().infer([row[1] for row in sample_data])
															a__ 							=	model.cuda()(
															    input_ids=hf_tokens['''input_ids'''].cuda()			,  attention_mask=hf_tokens['''attention_mask'''].cuda())
										else:
															a__ 							=	model(**lowercase_			,  output_hidden_states=lowercase_)
															a__ 							=	our_output['logits']
															if classification_head:
																				a__ 							=	esm.model.classification_heads['mnli'](esm.extract_features(lowercase_))
															else:
																				a__ 							=	esm(hf_tokens['''input_ids''']			,  repr_layers=list(range(999)))
																				a__ 							=	their_output['logits']

										if is_folding_model:
															a__ 							=	torch.max(torch.abs(our_output['''positions'''] - their_output['''positions'''])).item()
															a__ 							=	torch.allclose(our_output['''positions''']			,  their_output['''positions''']			,  atol=1E-5)
										else:
															a__ 							=	torch.max(torch.abs(our_output - their_output)).item()
															a__ 							=	torch.allclose(lowercase_			,  lowercase_			,  atol=1E-5)

										print(f'max_absolute_diff = {max_absolute_diff}')  # ~ 1e-5
										print('''Do both models output the same tensors?'''			,  '''🔥''' if success else '''💩''')

										if not success:
															raise Exception('''Something went wRoNg''')

										if not is_folding_model:
															# Let's check contact prediction too
															a__ 							=	model.predict_contacts(hf_tokens['''input_ids''']			,  hf_tokens['''attention_mask'''])
															a__ 							=	esm.predict_contacts(hf_tokens['''input_ids'''])
															a__ 							=	torch.max(torch.abs(our_output - their_output)).item()
															a__ 							=	torch.allclose(lowercase_			,  lowercase_			,  atol=1E-5)

															print('''Contact prediction testing:''')
															print(f'max_absolute_diff = {max_absolute_diff}')  # ~ 1e-5
															print('''Do both models output the same tensors?'''			,  '''🔥''' if success else '''💩''')

															if not success:
																				raise Exception('''Something went wRoNg''')

										pathlib.Path(lowercase_).mkdir(parents=lowercase_			,  exist_ok=lowercase_)
										print(f'Saving model to {pytorch_dump_folder_path}')
										model.save_pretrained(lowercase_)

										del esm  # Free up some memory before continuing

					print(f'Saving tokenizer to {pytorch_dump_folder_path}')
					hf_tokenizer.save_pretrained(lowercase_)

					if push_to_repo:
										model.push_to_hub(repo_id=lowercase_			,  use_auth_token=lowercase_)
										hf_tokenizer.push_to_hub(repo_id=lowercase_			,  use_auth_token=lowercase_)


if __name__ == "__main__":
	__a							:						str   							=			argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '--pytorch_dump_folder_path', type=str, required=True, help='Path to the output PyTorch model.'
	)
	parser.add_argument(
	    '--classification_head', action='store_true', help='Whether to convert a final classification head.'
	)
	parser.add_argument('--model', default=None, type=str, required=True, help='Name of model to convert.')
	parser.add_argument('--push_to_repo', type=str, help='Repo to upload to (including username!).')
	parser.add_argument('--auth_token', type=str, help='HuggingFace auth token.')
	__a							:						int   							=			parser.parse_args()
	convert_esm_checkpoint_to_pytorch(
	    args.model, args.pytorch_dump_folder_path, args.classification_head, args.push_to_repo, args.auth_token
	)



import unittest
from dataclasses import dataclass
from typing import List, Union

import numpy as np
import PIL.Image

from diffusers.utils.outputs import BaseOutput




@dataclass
class  __lowerCAmelCase	(							a_    ):



 '''simple docstring'''
 _lowerCamelCase    :						Union[List[PIL.Image.Image], np.ndarray]







class  __lowerCAmelCase	(							unittest.TestCase    ):



 '''simple docstring'''

 def        __snake_case       (      self							:		int	)			->				Tuple:
        _UpperCamelCase  :							str	    =	CustomOutput(images=np.random.rand(1							,      3							,      4							,      4	)	)

        # check every way of getting the attribute
        assert isinstance(outputs.images							,      np.ndarray	)
        assert outputs.images.shape == (1, 3, 4, 4)
        assert isinstance(outputs['images']							,      np.ndarray	)
        assert outputs["images"].shape == (1, 3, 4, 4)
        assert isinstance(outputs[0]							,      np.ndarray	)
        assert outputs[0].shape == (1, 3, 4, 4)

        # test with a non-tensor attribute
        _UpperCamelCase  :							str	    =	CustomOutput(images=[PIL.Image.new('RGB'							,      (4, 4)	)]	)

        # check every way of getting the attribute
        assert isinstance(outputs.images							,      lowerCamelCase	)
        assert isinstance(outputs.images[0]							,      PIL.Image.Image	)
        assert isinstance(outputs['images']							,      lowerCamelCase	)
        assert isinstance(outputs['images'][0]							,      PIL.Image.Image	)
        assert isinstance(outputs[0]							,      lowerCamelCase	)
        assert isinstance(outputs[0][0]							,      PIL.Image.Image	)

 def        __snake_case       (      self							:		Any	)			->				Dict:
        # test output reinitialization with a `dict` for compatibility with `accelerate`
        _UpperCamelCase  :							Tuple	    =	CustomOutput({'images': np.random.rand(1							,      3							,      4							,      4	)}	)

        # check every way of getting the attribute
        assert isinstance(outputs.images							,      np.ndarray	)
        assert outputs.images.shape == (1, 3, 4, 4)
        assert isinstance(outputs['images']							,      np.ndarray	)
        assert outputs["images"].shape == (1, 3, 4, 4)
        assert isinstance(outputs[0]							,      np.ndarray	)
        assert outputs[0].shape == (1, 3, 4, 4)

        # test with a non-tensor attribute
        _UpperCamelCase  :							int	    =	CustomOutput({'images': [PIL.Image.new('RGB'							,      (4, 4)	)]}	)

        # check every way of getting the attribute
        assert isinstance(outputs.images							,      lowerCamelCase	)
        assert isinstance(outputs.images[0]							,      PIL.Image.Image	)
        assert isinstance(outputs['images']							,      lowerCamelCase	)
        assert isinstance(outputs['images'][0]							,      PIL.Image.Image	)
        assert isinstance(outputs[0]							,      lowerCamelCase	)
        assert isinstance(outputs[0][0]							,      PIL.Image.Image	)



