code: 
  
  
  
  """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 :		str      							=    logging.get_logger(__name__)
  
  __A :		List[Any]      							=    [
      (
          '''protein1''',
          '''MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA''',
      ),
      ('''protein2''', '''MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA'''),
      ('''protein3''', '''MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG'''),
      ('''protein4''', '''MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA'''),
  ]
  
  __A :		Union[str, 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,
  }
  
  __A :		Union[str, Any]      							=    list('''ARNDCQEGHILKMFPSTWYV''')
  
  __A :		Optional[int]      							=    restypes + ['''X''']
  __A :		Any      							=    restypes_with_x + ['''<pad>''', '''<mask>''', '''<cls>''', '''<sep>''', '''<eos>''']
  def 					A_			(		):
  
  
  
  
  
  
  
      '''simple docstring'''
      with TemporaryDirectory() as tempdir:
          UpperCamelCase				:   Dict          =							"\n".join(UpperCamelCase__    )
          UpperCamelCase				:   Any          =							Path(UpperCamelCase__    ) / "vocab.txt"
          vocab_file.write_text(UpperCamelCase__    )
          UpperCamelCase				:   Union[str, Any]          =							EsmTokenizer(vocab_file=str(UpperCamelCase__    )    )
      UpperCamelCase				:   int          =							0  # Overlaps with 'A' but that seems to be what they want
      return hf_tokenizer
  def 					A_			(		snake_case_     :       List[Any]						,snake_case_     :       int    ):
  
  
  
  
  
  
  
      '''simple docstring'''
      UpperCamelCase				:   Tuple          =							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_			(		snake_case_     :       str						,snake_case_     :       str						,snake_case_     :       bool						,snake_case_     :       str						,snake_case_     :       str    ):
  
  
  
  
  
  
  
      '''simple docstring'''
      if model.startswith("""esmfold"""    ):
          UpperCamelCase				:   Dict          =							MODEL_MAPPING[model]()
      else:
          UpperCamelCase				:   List[Any]          =							MODEL_MAPPING[model]()
      esm.eval()  # disable dropout
  
      if model.startswith("""esmfold"""    ):
          UpperCamelCase				:   List[Any]          =							esm.esm.embed_dim
          UpperCamelCase				:   Tuple          =							esm.esm.num_layers
          UpperCamelCase				:   Optional[int]          =							esm.esm.attention_heads
          UpperCamelCase				:   Dict          =							4 * embed_dim
          UpperCamelCase				:   List[Any]          =							esm.esm.token_dropout
          UpperCamelCase				:   Any          =							False  # This code path does not exist in ESM-2
          UpperCamelCase				:   str          =							"rotary"
          UpperCamelCase				:   Optional[Any]          =							True
          UpperCamelCase				:   Any          =							EsmFoldConfig()
          for key, val in esm.cfg.items():
              if hasattr(UpperCamelCase__						,UpperCamelCase__    ) and key != "trunk":
                  setattr(UpperCamelCase__						,UpperCamelCase__						,UpperCamelCase__    )
          for key, val in esm.cfg.trunk.items():
              if hasattr(esmfold_config.trunk						,UpperCamelCase__    ) and key != "structure_module":
                  setattr(esmfold_config.trunk						,UpperCamelCase__						,UpperCamelCase__    )
          for key, val in esm.cfg.trunk.structure_module.items():
              if hasattr(esmfold_config.trunk.structure_module						,UpperCamelCase__    ):
                  setattr(esmfold_config.trunk.structure_module						,UpperCamelCase__						,UpperCamelCase__    )
      elif hasattr(UpperCamelCase__						,"""args"""    ):
          # Indicates an ESM-1b or ESM-1v model
          UpperCamelCase				:   List[Any]          =							esm.args.embed_dim
          UpperCamelCase				:   Optional[Any]          =							esm.args.layers
          UpperCamelCase				:   Union[str, Any]          =							esm.args.attention_heads
          UpperCamelCase				:   Dict          =							esm.args.ffn_embed_dim
          UpperCamelCase				:   str          =							esm.args.token_dropout
          UpperCamelCase				:   Dict          =							True if esm.emb_layer_norm_before else False
          UpperCamelCase				:   Optional[Any]          =							"absolute"
          UpperCamelCase				:   List[Any]          =							False
          UpperCamelCase				:   List[Any]          =							None
      else:
          # Indicates an ESM-2 model
          UpperCamelCase				:   Any          =							esm.embed_dim
          UpperCamelCase				:   Optional[int]          =							esm.num_layers
          UpperCamelCase				:   Union[str, Any]          =							esm.attention_heads
          UpperCamelCase				:   Any          =							4 * embed_dim  # This is hardcoded in ESM-2
          UpperCamelCase				:   Union[str, Any]          =							esm.token_dropout
          UpperCamelCase				:   List[str]          =							False  # This code path does not exist in ESM-2
          UpperCamelCase				:   Tuple          =							"rotary"
          UpperCamelCase				:   Tuple          =							False
          UpperCamelCase				:   Union[str, Any]          =							None
  
      if is_folding_model:
          UpperCamelCase				:   Dict          =							esm.esm.alphabet
      UpperCamelCase				:   Dict          =							tuple(alphabet.all_toks    )
      UpperCamelCase				:   List[str]          =							alphabet.mask_idx
      UpperCamelCase				:   Dict          =							alphabet.padding_idx
  
      if is_folding_model:
          UpperCamelCase				:   List[Any]          =							esm.esm
      else:
          UpperCamelCase				:   Union[str, Any]          =							esm
  
      UpperCamelCase				:   List[Any]          =							EsmConfig(
          vocab_size=original_esm_model.embed_tokens.num_embeddings						,mask_token_id=UpperCamelCase__						,hidden_size=UpperCamelCase__						,num_hidden_layers=UpperCamelCase__						,num_attention_heads=UpperCamelCase__						,intermediate_size=UpperCamelCase__						,max_position_embeddings=1_0_2_6						,layer_norm_eps=1e-5						,attention_probs_dropout_prob=0.0						,hidden_dropout_prob=0.0						,pad_token_id=UpperCamelCase__						,emb_layer_norm_before=UpperCamelCase__						,token_dropout=UpperCamelCase__						,position_embedding_type=UpperCamelCase__						,is_folding_model=UpperCamelCase__						,esmfold_config=UpperCamelCase__						,vocab_list=UpperCamelCase__						,)
      if classification_head:
          UpperCamelCase				:   int          =							esm.classification_heads["mnli"].out_proj.weight.shape[0]
      print("""Our ESM config:"""						,UpperCamelCase__    )
  
      if model.startswith("""esmfold"""    ):
          UpperCamelCase				:   str          =							EsmForProteinFolding
      elif classification_head:
          UpperCamelCase				:   Optional[int]          =							EsmForSequenceClassification
      else:
          UpperCamelCase				:   Tuple          =							EsmForMaskedLM
      UpperCamelCase				:   Any          =							model_class(UpperCamelCase__    )
      model.eval()
  
      # Now let's copy all the weights.
      # Embeddings
      UpperCamelCase				:   Optional[int]          =							original_esm_model.embed_tokens.weight
      if position_embedding_type == "absolute":
          UpperCamelCase				:   Dict          =							original_esm_model.embed_positions.weight
  
      if config.emb_layer_norm_before:
          UpperCamelCase				:   List[Any]          =							original_esm_model.emb_layer_norm_before.weight
          UpperCamelCase				:   Optional[int]          =							original_esm_model.emb_layer_norm_before.bias
  
      UpperCamelCase				:   int          =							original_esm_model.emb_layer_norm_after.weight
      UpperCamelCase				:   Tuple          =							original_esm_model.emb_layer_norm_after.bias
  
      for i in range(config.num_hidden_layers    ):
          # Encoder: start of layer
          UpperCamelCase				:   EsmLayer          =							model.esm.encoder.layer[i]
          # esm_layer: TransformerSentenceEncoderLayer = original_esm_model.layers[i]
          UpperCamelCase				:   List[Any]          =							original_esm_model.layers[i]
  
          # self attention
          UpperCamelCase				:   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)    )
          )
  
          UpperCamelCase				:   List[str]          =							esm_layer.self_attn.q_proj.weight
          UpperCamelCase				:   Optional[int]          =							esm_layer.self_attn.q_proj.bias
          UpperCamelCase				:   str          =							esm_layer.self_attn.k_proj.weight
          UpperCamelCase				:   List[Any]          =							esm_layer.self_attn.k_proj.bias
          UpperCamelCase				:   List[str]          =							esm_layer.self_attn.v_proj.weight
          UpperCamelCase				:   Union[str, Any]          =							esm_layer.self_attn.v_proj.bias
  
          if getattr(esm_layer.self_attn						,"""rot_emb"""						,UpperCamelCase__    ) 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				:   Union[str, Any]          =							esm_layer.self_attn.rot_emb.inv_freq
  
          # LayerNorm changes for pre-activation
          UpperCamelCase				:   Tuple          =							esm_layer.self_attn_layer_norm.weight
          UpperCamelCase				:   Optional[int]          =							esm_layer.self_attn_layer_norm.bias
          UpperCamelCase				:   Tuple          =							esm_layer.final_layer_norm.weight
          UpperCamelCase				:   str          =							esm_layer.final_layer_norm.bias
  
          # self-attention output
          UpperCamelCase				:   EsmSelfOutput          =							layer.attention.output
          assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
          UpperCamelCase				:   Tuple          =							esm_layer.self_attn.out_proj.weight
          UpperCamelCase				:   Any          =							esm_layer.self_attn.out_proj.bias
  
          # intermediate
          UpperCamelCase				:   EsmIntermediate          =							layer.intermediate
          assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
          UpperCamelCase				:   Union[str, Any]          =							esm_layer.fca.weight
          UpperCamelCase				:   Optional[int]          =							esm_layer.fca.bias
  
          # output
          UpperCamelCase				:   EsmOutput          =							layer.output
          assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
          UpperCamelCase				:   Union[str, Any]          =							esm_layer.fca.weight
          UpperCamelCase				:   Optional[Any]          =							esm_layer.fca.bias
          # end of layer
  
      if is_folding_model:
          UpperCamelCase				:   List[Any]          =							esm.esm_s_combine.data
          UpperCamelCase				:   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:
          UpperCamelCase				:   Union[str, Any]          =							esm.esm.classification_heads["mnli"].dense.weight
          UpperCamelCase				:   Dict          =							esm.classification_heads["mnli"].dense.bias
          UpperCamelCase				:   Any          =							esm.classification_heads["mnli"].out_proj.weight
          UpperCamelCase				:   Any          =							esm.classification_heads["mnli"].out_proj.bias
      else:
          # LM Head
          UpperCamelCase				:   Tuple          =							esm.lm_head.dense.weight
          UpperCamelCase				:   Any          =							esm.lm_head.dense.bias
          UpperCamelCase				:   List[str]          =							esm.lm_head.layer_norm.weight
          UpperCamelCase				:   str          =							esm.lm_head.layer_norm.bias
          UpperCamelCase				:   Any          =							esm.lm_head.weight
          UpperCamelCase				:   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.
          UpperCamelCase				:   Any          =							SAMPLE_DATA[:2]
      else:
          UpperCamelCase				:   Union[str, Any]          =							SAMPLE_DATA
  
      if is_folding_model:
          UpperCamelCase				:   Any          =							get_esmfold_tokenizer()
          UpperCamelCase				:   int          =							hf_tokenizer(
              [row[1] for row in sample_data]						,return_tensors="""pt"""						,padding=UpperCamelCase__						,add_special_tokens=UpperCamelCase__    )
          UpperCamelCase				:   Optional[Any]          =							esmfold_encode_sequences([row[1] for row in sample_data]    )
          UpperCamelCase				:   List[Any]          =							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				:   Any          =							alphabet.get_batch_converter()
          UpperCamelCase				:   List[str]          =							batch_converter(UpperCamelCase__    )
          # Prepare tokenizer and make sure it matches
          with TemporaryDirectory() as tempdir:
              UpperCamelCase				:   Any          =							"\n".join(alphabet.all_toks    )
              UpperCamelCase				:   str          =							Path(UpperCamelCase__    ) / "vocab.txt"
              vocab_file.write_text(UpperCamelCase__    )
              UpperCamelCase				:   Optional[Any]          =							EsmTokenizer(vocab_file=str(UpperCamelCase__    )    )
  
          UpperCamelCase				:   List[Any]          =							hf_tokenizer([row[1] for row in sample_data]						,return_tensors="""pt"""						,padding=UpperCamelCase__    )
          UpperCamelCase				:   List[Any]          =							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				:   Any          =							esm.cuda().infer([row[1] for row in sample_data]    )
              UpperCamelCase				:   str          =							model.cuda()(
                  input_ids=hf_tokens["""input_ids"""].cuda()						,attention_mask=hf_tokens["""attention_mask"""].cuda()    )
          else:
              UpperCamelCase				:   Union[str, Any]          =							model(**UpperCamelCase__						,output_hidden_states=UpperCamelCase__    )
              UpperCamelCase				:   Optional[Any]          =							our_output["logits"]
              if classification_head:
                  UpperCamelCase				:   List[str]          =							esm.model.classification_heads["mnli"](esm.extract_features(UpperCamelCase__    )    )
              else:
                  UpperCamelCase				:   int          =							esm(hf_tokens["""input_ids"""]						,repr_layers=list(range(9_9_9    )    )    )
                  UpperCamelCase				:   str          =							their_output["logits"]
  
          if is_folding_model:
              UpperCamelCase				:   Dict          =							torch.max(torch.abs(our_output["""positions"""] - their_output["""positions"""]    )    ).item()
              UpperCamelCase				:   List[str]          =							torch.allclose(our_output["""positions"""]						,their_output["""positions"""]						,atol=1e-5    )
          else:
              UpperCamelCase				:   List[Any]          =							torch.max(torch.abs(our_output - their_output    )    ).item()
              UpperCamelCase				:   List[Any]          =							torch.allclose(UpperCamelCase__						,UpperCamelCase__						,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				:   int          =							model.predict_contacts(hf_tokens["""input_ids"""]						,hf_tokens["""attention_mask"""]    )
              UpperCamelCase				:   Any          =							esm.predict_contacts(hf_tokens["""input_ids"""]    )
              UpperCamelCase				:   Optional[int]          =							torch.max(torch.abs(our_output - their_output    )    ).item()
              UpperCamelCase				:   Optional[Any]          =							torch.allclose(UpperCamelCase__						,UpperCamelCase__						,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(UpperCamelCase__    ).mkdir(parents=UpperCamelCase__						,exist_ok=UpperCamelCase__    )
          print(f'Saving model to {pytorch_dump_folder_path}'    )
          model.save_pretrained(UpperCamelCase__    )
  
          del esm  # Free up some memory before continuing
  
      print(f'Saving tokenizer to {pytorch_dump_folder_path}'    )
      hf_tokenizer.save_pretrained(UpperCamelCase__    )
  
      if push_to_repo:
          model.push_to_hub(repo_id=UpperCamelCase__						,use_auth_token=UpperCamelCase__    )
          hf_tokenizer.push_to_hub(repo_id=UpperCamelCase__						,use_auth_token=UpperCamelCase__    )
  
  
  if __name__ == "__main__":
    __A :		Optional[Any]      							=    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 :		Dict      							=    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
    )
  
  
  
  
  
style_context: 
  
  
  
  
  """simple docstring"""
  
  
  
  
  def 					A_			(		snake_case_     :       Tuple						,snake_case_     :       Optional[Any]						,snake_case_     :       Optional[Any]    ):
  
  
  
  
  
  
  
      '''simple docstring'''
      assert b != 0, "This cannot accept modulo that is == 0"
      if n == 0:
          return 1
  
      if n % 2 == 1:
          return (bin_exp_mod(snake_case_						,n - 1						,snake_case_    ) * a) % b
  
      UpperCamelCase				:   List[Any]          =							bin_exp_mod(snake_case_						,n / 2						,snake_case_    )
      return (r * r) % b
  
  
  if __name__ == "__main__":
    try:
      __A :		Any      							=    int(input('''Enter Base : ''').strip())
      __A :		Any      							=    int(input('''Enter Power : ''').strip())
      __A :		List[Any]      							=    int(input('''Enter Modulo : ''').strip())
    except ValueError:
      print('''Invalid literal for integer''')
  
    print(bin_exp_mod(BASE, POWER, MODULO))
  
  
  
  
  