




"""simple docstring"""





import json
import os
import shutil
import sys
import tempfile
import unittest
from unittest import TestCase
from unittest.mock import patch

import faiss
import numpy as np
from datasets import Dataset

from transformers import BartConfig, BartTokenizer, DPRConfig, DPRQuestionEncoderTokenizer, RagConfig
from transformers.file_utils import is_datasets_available, is_faiss_available, is_psutil_available, is_torch_available
from transformers.integrations import is_ray_available
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_ray


sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip

if is_torch_available():
							from distributed_pytorch_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
else:
							_lowercase	      =						None

if is_ray_available():
							import ray  # noqa: E402 # isort:skip
							from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
else:
							_lowercase	      =						None
							_lowercase	      =						None
							_lowercase	      =						None

def 		lowerCAmelCase							( lowerCamelCase		):
						'''simple docstring'''


						if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
												snake_case__       :					Tuple										=					unittest.skip('''test requires Datasets, Faiss, psutil'''		)(lowerCamelCase		)
						return test_case

@require_distributed_retrieval
class       __lowercase       (			snake_case_	):


				'''simple docstring'''


				def       _lowerCamelCase					(   self   :     Any      ):
										snake_case__       :					Union[str, Any]										=					tempfile.mkdtemp()
										snake_case__       :					str										=					8

										# DPR tok
										snake_case__       :					Tuple										=					[
										    '''[UNK]''',
										    '''[CLS]''',
										    '''[SEP]''',
										    '''[PAD]''',
										    '''[MASK]''',
										    '''want''',
										    '''##want''',
										    '''##ed''',
										    '''wa''',
										    '''un''',
										    '''runn''',
										    '''##ing''',
										    ''',''',
										    '''low''',
										    '''lowest''',
										]
										snake_case__       :					int										=					os.path.join(self.tmpdirname      ,'''dpr_tokenizer'''      )
										os.makedirs(SCREAMING_SNAKE_CASE_      ,exist_ok=SCREAMING_SNAKE_CASE_      )
										snake_case__       :					int										=					os.path.join(SCREAMING_SNAKE_CASE_      ,DPR_VOCAB_FILES_NAMES['''vocab_file''']      )
										with open(self.vocab_file      ,'''w'''      ,encoding='''utf-8'''      ) as vocab_writer:
																vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]      )      )

										# BART tok
										snake_case__       :					str										=					[
										    '''l''',
										    '''o''',
										    '''w''',
										    '''e''',
										    '''r''',
										    '''s''',
										    '''t''',
										    '''i''',
										    '''d''',
										    '''n''',
										    '''\u0120''',
										    '''\u0120l''',
										    '''\u0120n''',
										    '''\u0120lo''',
										    '''\u0120low''',
										    '''er''',
										    '''\u0120lowest''',
										    '''\u0120newer''',
										    '''\u0120wider''',
										    '''<unk>''',
										]
										snake_case__       :					Any										=					dict(zip(SCREAMING_SNAKE_CASE_      ,range(len(SCREAMING_SNAKE_CASE_      )      )      )      )
										snake_case__       :					List[str]										=					['''#version: 0.2''', '''\u0120 l''', '''\u0120l o''', '''\u0120lo w''', '''e r''', '''''']
										snake_case__       :					Dict										=					{'''unk_token''': '''<unk>'''}

										snake_case__       :					Optional[int]										=					os.path.join(self.tmpdirname      ,'''bart_tokenizer'''      )
										os.makedirs(SCREAMING_SNAKE_CASE_      ,exist_ok=SCREAMING_SNAKE_CASE_      )
										snake_case__       :					List[str]										=					os.path.join(SCREAMING_SNAKE_CASE_      ,BART_VOCAB_FILES_NAMES['''vocab_file''']      )
										snake_case__       :					Optional[Any]										=					os.path.join(SCREAMING_SNAKE_CASE_      ,BART_VOCAB_FILES_NAMES['''merges_file''']      )
										with open(self.vocab_file      ,'''w'''      ,encoding='''utf-8'''      ) as fp:
																fp.write(json.dumps(SCREAMING_SNAKE_CASE_      ) + '''\n'''      )
										with open(self.merges_file      ,'''w'''      ,encoding='''utf-8'''      ) as fp:
																fp.write('''\n'''.join(SCREAMING_SNAKE_CASE_      )      )

				def       _lowerCamelCase					(   self   :     List[Any]      ):
										return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname      ,'''dpr_tokenizer'''      )      )

				def       _lowerCamelCase					(   self   :     List[str]      ):
										return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname      ,'''bart_tokenizer'''      )      )

				def       _lowerCamelCase					(   self   :     Dict      ):
										shutil.rmtree(self.tmpdirname      )

				def       _lowerCamelCase					(   self   :     str      ):
										snake_case__       :					Dict										=					Dataset.from_dict(
										    {
										        '''id''': ['''0''', '''1'''],
										        '''text''': ['''foo''', '''bar'''],
										        '''title''': ['''Foo''', '''Bar'''],
										        '''embeddings''': [np.ones(self.retrieval_vector_size      ), 2 * np.ones(self.retrieval_vector_size      )],
										    }      )
										dataset.add_faiss_index('''embeddings'''      ,string_factory='''Flat'''      ,metric_type=faiss.METRIC_INNER_PRODUCT      )
										return dataset

				def       _lowerCamelCase					(   self   :     Dict      ,SCREAMING_SNAKE_CASE_   :     bool      ,SCREAMING_SNAKE_CASE_   :     str=1_2_3_4_5      ):
										snake_case__       :					List[str]										=					self.get_dummy_dataset()
										snake_case__       :					Dict										=					RagConfig(
										    retrieval_vector_size=self.retrieval_vector_size      ,question_encoder=DPRConfig().to_dict()      ,generator=BartConfig().to_dict()      ,)
										with patch('''transformers.models.rag.retrieval_rag.load_dataset'''      ) as mock_load_dataset:
																snake_case__       :					Optional[int]										=					dataset
																snake_case__       :					Optional[Any]										=					RagPyTorchDistributedRetriever(
																    SCREAMING_SNAKE_CASE_      ,question_encoder_tokenizer=self.get_dpr_tokenizer()      ,generator_tokenizer=self.get_bart_tokenizer()      ,)
																if init_retrieval:
																						retriever.init_retrieval(SCREAMING_SNAKE_CASE_      )
										return retriever

				def       _lowerCamelCase					(   self   :     Tuple      ,SCREAMING_SNAKE_CASE_   :     bool      ):
										# Have to run in local mode because sys.path modifications at top of
										# file are not propogated to remote workers.
										# https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
										ray.init(local_mode=SCREAMING_SNAKE_CASE_      )
										snake_case__       :					Dict										=					RagConfig(
										    retrieval_vector_size=self.retrieval_vector_size      ,question_encoder=DPRConfig().to_dict()      ,generator=BartConfig().to_dict()      ,)
										snake_case__       :					Any										=					ray.remote(SCREAMING_SNAKE_CASE_      )
										snake_case__       :					Dict										=					[remote_cls.remote() for _ in range(1      )]
										with patch('''transformers.models.rag.retrieval_rag.load_dataset'''      ) as mock_load_dataset:
																snake_case__       :					Dict										=					self.get_dummy_dataset()
																snake_case__       :					List[str]										=					RagRayDistributedRetriever(
																    SCREAMING_SNAKE_CASE_      ,question_encoder_tokenizer=self.get_dpr_tokenizer()      ,generator_tokenizer=self.get_bart_tokenizer()      ,retrieval_workers=SCREAMING_SNAKE_CASE_      ,)
																if init_retrieval:
																						retriever.init_retrieval()
										return retriever

				def       _lowerCamelCase					(   self   :     int      ,SCREAMING_SNAKE_CASE_   :     bool      ,SCREAMING_SNAKE_CASE_   :     bool      ,SCREAMING_SNAKE_CASE_   :     List[Any]=1_2_3_4_5      ):
										snake_case__       :					Optional[int]										=					self.get_dummy_dataset()
										snake_case__       :					Dict										=					RagConfig(
										    retrieval_vector_size=self.retrieval_vector_size      ,question_encoder=DPRConfig().to_dict()      ,generator=BartConfig().to_dict()      ,index_name='''custom'''      ,)
										if from_disk:
																snake_case__       :					Optional[int]										=					os.path.join(self.tmpdirname      ,'''dataset'''      )
																snake_case__       :					Optional[int]										=					os.path.join(self.tmpdirname      ,'''index.faiss'''      )
																dataset.get_index('''embeddings'''      ).save(os.path.join(self.tmpdirname      ,'''index.faiss'''      )      )
																dataset.drop_index('''embeddings'''      )
																dataset.save_to_disk(os.path.join(self.tmpdirname      ,'''dataset'''      )      )
																del dataset
																snake_case__       :					str										=					RagPyTorchDistributedRetriever(
																    SCREAMING_SNAKE_CASE_      ,question_encoder_tokenizer=self.get_dpr_tokenizer()      ,generator_tokenizer=self.get_bart_tokenizer()      ,)
										else:
																snake_case__       :					str										=					RagPyTorchDistributedRetriever(
																    SCREAMING_SNAKE_CASE_      ,question_encoder_tokenizer=self.get_dpr_tokenizer()      ,generator_tokenizer=self.get_bart_tokenizer()      ,index=CustomHFIndex(config.retrieval_vector_size      ,SCREAMING_SNAKE_CASE_      )      ,)
										if init_retrieval:
																retriever.init_retrieval(SCREAMING_SNAKE_CASE_      )
										return retriever

				def       _lowerCamelCase					(   self   :     List[Any]      ,SCREAMING_SNAKE_CASE_   :     bool      ,SCREAMING_SNAKE_CASE_   :     bool      ):
										# Have to run in local mode because sys.path modifications at top of
										# file are not propogated to remote workers.
										# https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
										ray.init(local_mode=SCREAMING_SNAKE_CASE_      )
										snake_case__       :					Optional[Any]										=					self.get_dummy_dataset()
										snake_case__       :					List[Any]										=					RagConfig(
										    retrieval_vector_size=self.retrieval_vector_size      ,question_encoder=DPRConfig().to_dict()      ,generator=BartConfig().to_dict()      ,index_name='''custom'''      ,)
										snake_case__       :					Optional[int]										=					ray.remote(SCREAMING_SNAKE_CASE_      )
										snake_case__       :					Optional[int]										=					[remote_cls.remote() for _ in range(1      )]
										if from_disk:
																snake_case__       :					str										=					os.path.join(self.tmpdirname      ,'''dataset'''      )
																snake_case__       :					List[Any]										=					os.path.join(self.tmpdirname      ,'''index.faiss'''      )
																dataset.get_index('''embeddings'''      ).save(os.path.join(self.tmpdirname      ,'''index.faiss'''      )      )
																dataset.drop_index('''embeddings'''      )
																dataset.save_to_disk(os.path.join(self.tmpdirname      ,'''dataset'''      )      )
																del dataset
																snake_case__       :					Any										=					RagRayDistributedRetriever(
																    SCREAMING_SNAKE_CASE_      ,question_encoder_tokenizer=self.get_dpr_tokenizer()      ,generator_tokenizer=self.get_bart_tokenizer()      ,retrieval_workers=SCREAMING_SNAKE_CASE_      ,index=CustomHFIndex.load_from_disk(
																        vector_size=config.retrieval_vector_size      ,dataset_path=config.passages_path      ,index_path=config.index_path      ,)      ,)
										else:
																snake_case__       :					Union[str, Any]										=					RagRayDistributedRetriever(
																    SCREAMING_SNAKE_CASE_      ,question_encoder_tokenizer=self.get_dpr_tokenizer()      ,generator_tokenizer=self.get_bart_tokenizer()      ,retrieval_workers=SCREAMING_SNAKE_CASE_      ,index=CustomHFIndex(config.retrieval_vector_size      ,SCREAMING_SNAKE_CASE_      )      ,)
										if init_retrieval:
																retriever.init_retrieval()
										return retriever

				def       _lowerCamelCase					(   self   :     List[Any]      ,SCREAMING_SNAKE_CASE_   :     RagRetriever      ,SCREAMING_SNAKE_CASE_   :     np.array      ,SCREAMING_SNAKE_CASE_   :     int      ):
										snake_case__       :					str										=					retriever.retrieve(SCREAMING_SNAKE_CASE_      ,n_docs=SCREAMING_SNAKE_CASE_      )
										self.assertEqual(retrieved_doc_embeds.shape      ,(2, n_docs, self.retrieval_vector_size)      )
										self.assertEqual(len(SCREAMING_SNAKE_CASE_      )      ,2      )
										self.assertEqual(sorted(doc_dicts[0]      )      ,['''embeddings''', '''id''', '''text''', '''title''']      )
										self.assertEqual(len(doc_dicts[0]['''id''']      )      ,SCREAMING_SNAKE_CASE_      )
										self.assertEqual(doc_dicts[0]['''id'''][0]      ,'''1'''      )  # max inner product is reached with second doc
										self.assertEqual(doc_dicts[1]['''id'''][0]      ,'''0'''      )  # max inner product is reached with first doc
										self.assertListEqual(doc_ids.tolist()      ,[[1], [0]]      )

				def       _lowerCamelCase					(   self   :     List[str]      ):
										snake_case__       :					Dict										=					1
										snake_case__       :					List[Any]										=					np.array(
										    [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]      ,dtype=np.floataa      )

										self.distributed_retriever_check(
										    self.get_dummy_pytorch_distributed_retriever(init_retrieval=SCREAMING_SNAKE_CASE_      )      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      )

				def       _lowerCamelCase					(   self   :     List[Any]      ):
										snake_case__       :					Any										=					1
										snake_case__       :					Dict										=					np.array(
										    [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]      ,dtype=np.floataa      )

										self.distributed_retriever_check(
										    self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=SCREAMING_SNAKE_CASE_      ,from_disk=SCREAMING_SNAKE_CASE_      )      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,)

				def       _lowerCamelCase					(   self   :     Optional[Any]      ):
										snake_case__       :					Union[str, Any]										=					1
										snake_case__       :					Tuple										=					np.array(
										    [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]      ,dtype=np.floataa      )

										self.distributed_retriever_check(
										    self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=SCREAMING_SNAKE_CASE_      ,from_disk=SCREAMING_SNAKE_CASE_      )      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,)

				@require_ray
				def       _lowerCamelCase					(   self   :     List[Any]      ):
										snake_case__       :					Optional[int]										=					1
										snake_case__       :					Optional[int]										=					np.array(
										    [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]      ,dtype=np.floataa      )

										self.distributed_retriever_check(
										    self.get_dummy_ray_distributed_retriever(init_retrieval=SCREAMING_SNAKE_CASE_      )      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      )
										ray.shutdown()

				@require_ray
				def       _lowerCamelCase					(   self   :     List[str]      ):
										snake_case__       :					Tuple										=					1
										snake_case__       :					Tuple										=					np.array(
										    [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]      ,dtype=np.floataa      )
										with self.assertRaises(SCREAMING_SNAKE_CASE_      ):
																self.distributed_retriever_check(
																    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=SCREAMING_SNAKE_CASE_      ,from_disk=SCREAMING_SNAKE_CASE_      )      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,)
										ray.shutdown()





				@require_ray
				def       _lowerCamelCase					(   self   :     Dict      ):
										snake_case__       :					List[Any]										=					1
										snake_case__       :					str										=					np.array(
										    [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )]      ,dtype=np.floataa      )

										self.distributed_retriever_check(
										    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=SCREAMING_SNAKE_CASE_      ,from_disk=SCREAMING_SNAKE_CASE_      )      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      )
										ray.shutdown()








"""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()
_lowercase	      =    logging.get_logger(__name__)

_lowercase	      =    [
    (
        """protein1""",
        """MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA""",
    ),
    ("""protein2""", """MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA"""),
    ("""protein3""", """MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG"""),
    ("""protein4""", """MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA"""),
]

_lowercase	      =    {
    """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,
}

_lowercase	      =    list("""ARNDCQEGHILKMFPSTWYV""")

_lowercase	      =    restypes + ["""X"""]
_lowercase	      =    restypes_with_x + ["""<pad>""", """<mask>""", """<cls>""", """<sep>""", """<eos>"""]

def 		lowerCAmelCase							( ):
						'''simple docstring'''


						with TemporaryDirectory() as tempdir:
												snake_case__       :					Any										=					'''\n'''.join(lowerCamelCase		)
												snake_case__       :					int										=					Path(lowerCamelCase		) / '''vocab.txt'''
												vocab_file.write_text(lowerCamelCase		)
												snake_case__       :					List[str]										=					EsmTokenizer(vocab_file=str(lowerCamelCase		)		)
						snake_case__       :					Dict										=					0  # Overlaps with 'A' but that seems to be what they want
						return hf_tokenizer

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


						snake_case__       :					Union[str, Any]										=					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 		lowerCAmelCase							( lowerCamelCase      ,			lowerCamelCase      ,			lowerCamelCase      ,			lowerCamelCase      ,			lowerCamelCase		):
						'''simple docstring'''


						if model.startswith('''esmfold'''		):
												snake_case__       :					Dict										=					MODEL_MAPPING[model]()
						else:
												snake_case__			,	snake_case__       :					List[str]										=					MODEL_MAPPING[model]()
						esm.eval()  # disable dropout

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

						if is_folding_model:
												snake_case__       :					int										=					esm.esm.alphabet
						snake_case__       :					Tuple										=					tuple(alphabet.all_toks		)
						snake_case__       :					List[str]										=					alphabet.mask_idx
						snake_case__       :					str										=					alphabet.padding_idx

						if is_folding_model:
												snake_case__       :					Any										=					esm.esm
						else:
												snake_case__       :					List[str]										=					esm

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

						if model.startswith('''esmfold'''		):
												snake_case__       :					List[str]										=					EsmForProteinFolding
						elif classification_head:
												snake_case__       :					str										=					EsmForSequenceClassification
						else:
												snake_case__       :					Dict										=					EsmForMaskedLM
						snake_case__       :					str										=					model_class(lowerCamelCase		)
						model.eval()

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

						if config.emb_layer_norm_before:
												snake_case__       :					Union[str, Any]										=					original_esm_model.emb_layer_norm_before.weight
												snake_case__       :					Union[str, Any]										=					original_esm_model.emb_layer_norm_before.bias

						snake_case__       :					str										=					original_esm_model.emb_layer_norm_after.weight
						snake_case__       :					Optional[int]										=					original_esm_model.emb_layer_norm_after.bias

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

												# self attention
												snake_case__       :					EsmSelfAttention										=					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)		)
												)

												snake_case__       :					Union[str, Any]										=					esm_layer.self_attn.q_proj.weight
												snake_case__       :					int										=					esm_layer.self_attn.q_proj.bias
												snake_case__       :					Tuple										=					esm_layer.self_attn.k_proj.weight
												snake_case__       :					Union[str, Any]										=					esm_layer.self_attn.k_proj.bias
												snake_case__       :					List[str]										=					esm_layer.self_attn.v_proj.weight
												snake_case__       :					Union[str, Any]										=					esm_layer.self_attn.v_proj.bias

												if getattr(esm_layer.self_attn      ,			'''rot_emb'''      ,			lowerCamelCase		) 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.
																		snake_case__       :					Optional[Any]										=					esm_layer.self_attn.rot_emb.inv_freq

												# LayerNorm changes for pre-activation
												snake_case__       :					int										=					esm_layer.self_attn_layer_norm.weight
												snake_case__       :					Tuple										=					esm_layer.self_attn_layer_norm.bias
												snake_case__       :					Union[str, Any]										=					esm_layer.final_layer_norm.weight
												snake_case__       :					Union[str, Any]										=					esm_layer.final_layer_norm.bias

												# self-attention output
												snake_case__       :					EsmSelfOutput										=					layer.attention.output
												assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
												snake_case__       :					List[Any]										=					esm_layer.self_attn.out_proj.weight
												snake_case__       :					Dict										=					esm_layer.self_attn.out_proj.bias

												# intermediate
												snake_case__       :					EsmIntermediate										=					layer.intermediate
												assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
												snake_case__       :					Dict										=					esm_layer.fca.weight
												snake_case__       :					Optional[int]										=					esm_layer.fca.bias

												# output
												snake_case__       :					EsmOutput										=					layer.output
												assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
												snake_case__       :					List[str]										=					esm_layer.fca.weight
												snake_case__       :					Optional[Any]										=					esm_layer.fca.bias
												# end of layer

						if is_folding_model:
												snake_case__       :					List[str]										=					esm.esm_s_combine.data
												snake_case__       :					Tuple										=					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:
												snake_case__       :					str										=					esm.esm.classification_heads['''mnli'''].dense.weight
												snake_case__       :					Optional[Any]										=					esm.classification_heads['''mnli'''].dense.bias
												snake_case__       :					Dict										=					esm.classification_heads['''mnli'''].out_proj.weight
												snake_case__       :					int										=					esm.classification_heads['''mnli'''].out_proj.bias
						else:
												# LM Head
												snake_case__       :					Optional[Any]										=					esm.lm_head.dense.weight
												snake_case__       :					Dict										=					esm.lm_head.dense.bias
												snake_case__       :					Tuple										=					esm.lm_head.layer_norm.weight
												snake_case__       :					List[Any]										=					esm.lm_head.layer_norm.bias
												snake_case__       :					List[str]										=					esm.lm_head.weight
												snake_case__       :					Optional[int]										=					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.
												snake_case__       :					int										=					SAMPLE_DATA[:2]
						else:
												snake_case__       :					Tuple										=					SAMPLE_DATA

						if is_folding_model:
												snake_case__       :					Dict										=					get_esmfold_tokenizer()
												snake_case__       :					Optional[int]										=					hf_tokenizer(
												    [row[1] for row in sample_data]      ,			return_tensors='''pt'''      ,			padding=lowerCamelCase      ,			add_special_tokens=lowerCamelCase		)
												snake_case__			,	snake_case__			,	snake_case__			,	snake_case__			,	snake_case__       :					Union[str, Any]										=					esmfold_encode_sequences([row[1] for row in sample_data]		)
												snake_case__       :					List[str]										=					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.
												snake_case__       :					Dict										=					alphabet.get_batch_converter()
												snake_case__			,	snake_case__			,	snake_case__       :					Tuple										=					batch_converter(lowerCamelCase		)
												# Prepare tokenizer and make sure it matches
												with TemporaryDirectory() as tempdir:
																		snake_case__       :					Tuple										=					'''\n'''.join(alphabet.all_toks		)
																		snake_case__       :					Union[str, Any]										=					Path(lowerCamelCase		) / '''vocab.txt'''
																		vocab_file.write_text(lowerCamelCase		)
																		snake_case__       :					List[Any]										=					EsmTokenizer(vocab_file=str(lowerCamelCase		)		)

												snake_case__       :					Optional[Any]										=					hf_tokenizer([row[1] for row in sample_data]      ,			return_tensors='''pt'''      ,			padding=lowerCamelCase		)
												snake_case__       :					Dict										=					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.
																		snake_case__       :					Optional[int]										=					esm.cuda().infer([row[1] for row in sample_data]		)
																		snake_case__       :					Union[str, Any]										=					model.cuda()(
																		    input_ids=hf_tokens['''input_ids'''].cuda()      ,			attention_mask=hf_tokens['''attention_mask'''].cuda()		)
												else:
																		snake_case__       :					List[str]										=					model(**lowerCamelCase      ,			output_hidden_states=lowerCamelCase		)
																		snake_case__       :					List[Any]										=					our_output['''logits''']
																		if classification_head:
																								snake_case__       :					Optional[Any]										=					esm.model.classification_heads['''mnli'''](esm.extract_features(lowerCamelCase		)		)
																		else:
																								snake_case__       :					Any										=					esm(hf_tokens['''input_ids''']      ,			repr_layers=list(range(999		)		)		)
																								snake_case__       :					Union[str, Any]										=					their_output['''logits''']

												if is_folding_model:
																		snake_case__       :					str										=					torch.max(torch.abs(our_output['''positions'''] - their_output['''positions''']		)		).item()
																		snake_case__       :					List[str]										=					torch.allclose(our_output['''positions''']      ,			their_output['''positions''']      ,			atol=1e-5		)
												else:
																		snake_case__       :					Tuple										=					torch.max(torch.abs(our_output - their_output		)		).item()
																		snake_case__       :					List[Any]										=					torch.allclose(lowerCamelCase      ,			lowerCamelCase      ,			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
																		snake_case__       :					str										=					model.predict_contacts(hf_tokens['''input_ids''']      ,			hf_tokens['''attention_mask''']		)
																		snake_case__       :					int										=					esm.predict_contacts(hf_tokens['''input_ids''']		)
																		snake_case__       :					Any										=					torch.max(torch.abs(our_output - their_output		)		).item()
																		snake_case__       :					Optional[Any]										=					torch.allclose(lowerCamelCase      ,			lowerCamelCase      ,			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(lowerCamelCase		).mkdir(parents=lowerCamelCase      ,			exist_ok=lowerCamelCase		)
												print(f'''Saving model to {pytorch_dump_folder_path}'''		)
												model.save_pretrained(lowerCamelCase		)

												del esm  # Free up some memory before continuing

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

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


if __name__ == "__main__":
							_lowercase	      =    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.""")
							_lowercase	      =    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
							)


