





'''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()
UpperCAmelCase       :		Tuple		=				logging.get_logger(__name__)

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

UpperCAmelCase       :		List[Any]		=				{
    '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,
}

UpperCAmelCase       :		int		=				list('ARNDCQEGHILKMFPSTWYV')

UpperCAmelCase       :		str		=				restypes + ['X']
UpperCAmelCase       :		Any		=				restypes_with_x + ['<pad>', '<mask>', '<cls>', '<sep>', '<eos>']




def  _a     (				):





      """simple docstring"""



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




def  _a     (				lowerCAmelCase_      ,   lowerCAmelCase_						):





      """simple docstring"""



      _snake_case		:						str         =  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  _a     (				lowerCAmelCase_      ,   lowerCAmelCase_      ,   lowerCAmelCase_      ,   lowerCAmelCase_      ,   lowerCAmelCase_						):





      """simple docstring"""



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

      if model.startswith('''esmfold'''						):
            _snake_case		:						List[str]         =  esm.esm.embed_dim
            _snake_case		:						Tuple         =  esm.esm.num_layers
            _snake_case		:						Dict         =  esm.esm.attention_heads
            _snake_case		:						Any         =  4 * embed_dim
            _snake_case		:						Union[str, Any]         =  esm.esm.token_dropout
            _snake_case		:						Optional[Any]         =  False  # This code path does not exist in ESM-2
            _snake_case		:						Optional[Any]         =  '''rotary'''
            _snake_case		:						Tuple         =  True
            _snake_case		:						List[Any]         =  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		:						Any         =  esm.args.embed_dim
            _snake_case		:						Optional[int]         =  esm.args.layers
            _snake_case		:						List[Any]         =  esm.args.attention_heads
            _snake_case		:						Optional[Any]         =  esm.args.ffn_embed_dim
            _snake_case		:						List[Any]         =  esm.args.token_dropout
            _snake_case		:						str         =  True if esm.emb_layer_norm_before else False
            _snake_case		:						int         =  '''absolute'''
            _snake_case		:						Tuple         =  False
            _snake_case		:						List[str]         =  None
      else:
            # Indicates an ESM-2 model
            _snake_case		:						Tuple         =  esm.embed_dim
            _snake_case		:						Tuple         =  esm.num_layers
            _snake_case		:						Dict         =  esm.attention_heads
            _snake_case		:						int         =  4 * embed_dim  # This is hardcoded in ESM-2
            _snake_case		:						Dict         =  esm.token_dropout
            _snake_case		:						Optional[int]         =  False  # This code path does not exist in ESM-2
            _snake_case		:						Dict         =  '''rotary'''
            _snake_case		:						Dict         =  False
            _snake_case		:						Dict         =  None

      if is_folding_model:
            _snake_case		:						List[str]         =  esm.esm.alphabet
      _snake_case		:						List[Any]         =  tuple(alphabet.all_toks						)
      _snake_case		:						str         =  alphabet.mask_idx
      _snake_case		:						Any         =  alphabet.padding_idx

      if is_folding_model:
            _snake_case		:						Tuple         =  esm.esm
      else:
            _snake_case		:						Dict         =  esm

      _snake_case		:						int         =  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=1_026      ,   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		:						List[Any]         =  esm.classification_heads['''mnli'''].out_proj.weight.shape[0]
      print('''Our ESM config:'''      ,   lowerCAmelCase_						)

      if model.startswith('''esmfold'''						):
            _snake_case		:						List[Any]         =  EsmForProteinFolding
      elif classification_head:
            _snake_case		:						Dict         =  EsmForSequenceClassification
      else:
            _snake_case		:						Optional[Any]         =  EsmForMaskedLM
      _snake_case		:						Optional[int]         =  model_class(lowerCAmelCase_						)
      model.eval()

      # Now let's copy all the weights.
      # Embeddings
      _snake_case		:						List[Any]         =  original_esm_model.embed_tokens.weight
      if position_embedding_type == "absolute":
            _snake_case		:						Optional[Any]         =  original_esm_model.embed_positions.weight

      if config.emb_layer_norm_before:
            _snake_case		:						Tuple         =  original_esm_model.emb_layer_norm_before.weight
            _snake_case		:						List[Any]         =  original_esm_model.emb_layer_norm_before.bias

      _snake_case		:						str         =  original_esm_model.emb_layer_norm_after.weight
      _snake_case		:						Dict         =  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		:						int         =  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		:						Optional[Any]         =  esm_layer.self_attn.q_proj.weight
            _snake_case		:						List[Any]         =  esm_layer.self_attn.q_proj.bias
            _snake_case		:						List[Any]         =  esm_layer.self_attn.k_proj.weight
            _snake_case		:						Dict         =  esm_layer.self_attn.k_proj.bias
            _snake_case		:						Union[str, Any]         =  esm_layer.self_attn.v_proj.weight
            _snake_case		:						List[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[int]         =  esm_layer.self_attn.rot_emb.inv_freq

            # LayerNorm changes for pre-activation
            _snake_case		:						Optional[int]         =  esm_layer.self_attn_layer_norm.weight
            _snake_case		:						int         =  esm_layer.self_attn_layer_norm.bias
            _snake_case		:						str         =  esm_layer.final_layer_norm.weight
            _snake_case		:						str         =  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		:						Optional[int]         =  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		:						Tuple         =  esm_layer.fca.weight
            _snake_case		:						str         =  esm_layer.fca.bias

            # output
            _snake_case		:						EsmOutput         =  layer.output
            assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
            _snake_case		:						Tuple         =  esm_layer.fca.weight
            _snake_case		:						Dict         =  esm_layer.fca.bias
            # end of layer

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

      if is_folding_model:
            _snake_case		:						Any         =  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		:						int         =  esmfold_encode_sequences([row[1] for row in sample_data]						)
            _snake_case		:						int         =  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		:						List[str]         =  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		:						int         =  '''\n'''.join(alphabet.all_toks						)
                  _snake_case		:						List[Any]         =  Path(lowerCAmelCase_						) / '''vocab.txt'''
                  vocab_file.write_text(lowerCAmelCase_						)
                  _snake_case		:						int         =  EsmTokenizer(vocab_file=str(lowerCAmelCase_						)						)

            _snake_case		:						List[str]         =  hf_tokenizer([row[1] for row in sample_data]      ,   return_tensors='''pt'''      ,   padding=lowerCAmelCase_						)
            _snake_case		:						str         =  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		:						Optional[int]         =  model.cuda()(
                      input_ids=hf_tokens['''input_ids'''].cuda()      ,   attention_mask=hf_tokens['''attention_mask'''].cuda()						)
            else:
                  _snake_case		:						Tuple         =  model(**lowerCAmelCase_      ,   output_hidden_states=lowerCAmelCase_						)
                  _snake_case		:						int         =  our_output['''logits''']
                  if classification_head:
                        _snake_case		:						Tuple         =  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		:						Optional[Any]         =  their_output['''logits''']

            if is_folding_model:
                  _snake_case		:						Any         =  torch.max(torch.abs(our_output['''positions'''] - their_output['''positions''']						)						).item()
                  _snake_case		:						Tuple         =  torch.allclose(our_output['''positions''']      ,   their_output['''positions''']      ,   atol=1E-5						)
            else:
                  _snake_case		:						str         =  torch.max(torch.abs(our_output - their_output						)						).item()
                  _snake_case		:						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		:						Any         =  esm.predict_contacts(hf_tokens['''input_ids''']						)
                  _snake_case		:						Any         =  torch.max(torch.abs(our_output - their_output						)						).item()
                  _snake_case		:						Union[str, 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__":
   UpperCAmelCase       :		int		=				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.')
   UpperCAmelCase       :		Optional[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
   )









'''simple docstring'''

from __future__ import annotations




def  _a     (				lowerCAmelCase_      ,   lowerCAmelCase_      ,   lowerCAmelCase_      ,   lowerCAmelCase_						):





      """simple docstring"""



      if counter == 0 or capacity == 0:
            return 0

      # If weight of the nth item is more than Knapsack of capacity,
      #   then this item cannot be included in the optimal solution,
      # else return the maximum of two cases:
      #   (1) nth item included
      #   (2) not included
      if weights[counter - 1] > capacity:
            return knapsack(lowerCAmelCase_      ,   lowerCAmelCase_      ,   lowerCAmelCase_      ,   counter - 1						)
      else:
            _snake_case		:						List[str]         =  capacity - weights[counter - 1]
            _snake_case		:						Tuple         =  values[counter - 1] + knapsack(
                lowerCAmelCase_      ,   lowerCAmelCase_      ,   lowerCAmelCase_      ,   counter - 1						)
            _snake_case		:						Any         =  knapsack(lowerCAmelCase_      ,   lowerCAmelCase_      ,   lowerCAmelCase_      ,   counter - 1						)
            return max(lowerCAmelCase_      ,   lowerCAmelCase_						)


if __name__ == "__main__":
   import doctest

   doctest.testmod()


