


'''simple docstring'''







from __future__ import annotations

import sys







def _UpperCamelCase			(	__A ,     __A       )				->						tuple[int, int]:


      '''simple docstring'''



      if abs(__A       ) == 1:
            return a, 0
      elif abs(__A       ) == 1:
            return 0, b

      UpperCamelCase__   ,      UpperCamelCase__      =							a, b
      UpperCamelCase__   ,      UpperCamelCase__      =							1, 0
      UpperCamelCase__   ,      UpperCamelCase__      =							0, 1

      while remainder != 0:
            UpperCamelCase__      =							old_remainder // remainder
            UpperCamelCase__   ,      UpperCamelCase__      =							remainder, old_remainder - quotient * remainder
            UpperCamelCase__   ,      UpperCamelCase__      =							coeff_a, old_coeff_a - quotient * coeff_a
            UpperCamelCase__   ,      UpperCamelCase__      =							coeff_b, old_coeff_b - quotient * coeff_b

      # sign correction for negative numbers
      if a < 0:
            UpperCamelCase__      =							-old_coeff_a
      if b < 0:
            UpperCamelCase__      =							-old_coeff_b

      return old_coeff_a, old_coeff_b







def _UpperCamelCase			(	)				->						Dict:


      '''simple docstring'''



      if len(sys.argv       ) < 3:
            print("2 integer arguments required"       )
            return 1
      UpperCamelCase__      =							int(sys.argv[1]       )
      UpperCamelCase__      =							int(sys.argv[2]       )
      print(extended_euclidean_algorithm(__A ,     __A       )       )
      return 0


if __name__ == "__main__":
  raise SystemExit(main())









'''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__    :   Optional[int]								=			logging.get_logger(__name__)

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

a__    :   Optional[int]								=			{
    '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__    :   List[Any]								=			list('ARNDCQEGHILKMFPSTWYV')

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







def _UpperCamelCase			(	)				->						str:


      '''simple docstring'''



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







def _UpperCamelCase			(	__A ,     __A       )				->						List[str]:


      '''simple docstring'''



      UpperCamelCase__      =							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 _UpperCamelCase			(	__A ,     __A ,     __A ,     __A ,     __A       )				->						int:


      '''simple docstring'''



      if model.startswith("esmfold"       ):
            UpperCamelCase__      =							MODEL_MAPPING[model]()
      else:
            UpperCamelCase__   ,      UpperCamelCase__      =							MODEL_MAPPING[model]()
      esm.eval()  # disable dropout

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

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

      if is_folding_model:
            UpperCamelCase__      =							esm.esm
      else:
            UpperCamelCase__      =							esm

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

      if model.startswith("esmfold"       ):
            UpperCamelCase__      =							EsmForProteinFolding
      elif classification_head:
            UpperCamelCase__      =							EsmForSequenceClassification
      else:
            UpperCamelCase__      =							EsmForMaskedLM
      UpperCamelCase__      =							model_class(__A       )
      model.eval()

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

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

      UpperCamelCase__      =							original_esm_model.emb_layer_norm_after.weight
      UpperCamelCase__      =							original_esm_model.emb_layer_norm_after.bias

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

            # self attention
            UpperCamelCase__      =							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)       )
            )

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

            if getattr(esm_layer.self_attn ,     "rot_emb" ,     __A       ) 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.
                  UpperCamelCase__      =							esm_layer.self_attn.rot_emb.inv_freq

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

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

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

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

      if is_folding_model:
            UpperCamelCase__      =							esm.esm_s_combine.data
            UpperCamelCase__      =							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:
            UpperCamelCase__      =							esm.esm.classification_heads["mnli"].dense.weight
            UpperCamelCase__      =							esm.classification_heads["mnli"].dense.bias
            UpperCamelCase__      =							esm.classification_heads["mnli"].out_proj.weight
            UpperCamelCase__      =							esm.classification_heads["mnli"].out_proj.bias
      else:
            # LM Head
            UpperCamelCase__      =							esm.lm_head.dense.weight
            UpperCamelCase__      =							esm.lm_head.dense.bias
            UpperCamelCase__      =							esm.lm_head.layer_norm.weight
            UpperCamelCase__      =							esm.lm_head.layer_norm.bias
            UpperCamelCase__      =							esm.lm_head.weight
            UpperCamelCase__      =							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.
            UpperCamelCase__      =							SAMPLE_DATA[:2]
      else:
            UpperCamelCase__      =							SAMPLE_DATA

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

            UpperCamelCase__      =							hf_tokenizer([row[1] for row in sample_data] ,     return_tensors="pt" ,     padding=__A       )
            UpperCamelCase__      =							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.
                  UpperCamelCase__      =							esm.cuda().infer([row[1] for row in sample_data]       )
                  UpperCamelCase__      =							model.cuda()(
                      input_ids=hf_tokens["input_ids"].cuda() ,     attention_mask=hf_tokens["attention_mask"].cuda()       )
            else:
                  UpperCamelCase__      =							model(**__A ,     output_hidden_states=__A       )
                  UpperCamelCase__      =							our_output["logits"]
                  if classification_head:
                        UpperCamelCase__      =							esm.model.classification_heads["mnli"](esm.extract_features(__A       )       )
                  else:
                        UpperCamelCase__      =							esm(hf_tokens["input_ids"] ,     repr_layers=list(range(999       )       )       )
                        UpperCamelCase__      =							their_output["logits"]

            if is_folding_model:
                  UpperCamelCase__      =							torch.max(torch.abs(our_output["positions"] - their_output["positions"]       )       ).item()
                  UpperCamelCase__      =							torch.allclose(our_output["positions"] ,     their_output["positions"] ,     atol=1E-5       )
            else:
                  UpperCamelCase__      =							torch.max(torch.abs(our_output - their_output       )       ).item()
                  UpperCamelCase__      =							torch.allclose(__A ,     __A ,     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
                  UpperCamelCase__      =							model.predict_contacts(hf_tokens["input_ids"] ,     hf_tokens["attention_mask"]       )
                  UpperCamelCase__      =							esm.predict_contacts(hf_tokens["input_ids"]       )
                  UpperCamelCase__      =							torch.max(torch.abs(our_output - their_output       )       ).item()
                  UpperCamelCase__      =							torch.allclose(__A ,     __A ,     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(__A       ).mkdir(parents=__A ,     exist_ok=__A       )
            print(F'''Saving model to {pytorch_dump_folder_path}'''       )
            model.save_pretrained(__A       )

            del esm  # Free up some memory before continuing

      print(F'''Saving tokenizer to {pytorch_dump_folder_path}'''       )
      hf_tokenizer.save_pretrained(__A       )

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


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__    :   Optional[Any]								=			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
  )





