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

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

__UpperCAmelCase												=   {
    '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												=   list('ARNDCQEGHILKMFPSTWYV')

__UpperCAmelCase												=   restypes + ['X']
__UpperCAmelCase												=   restypes_with_x + ['<pad>', '<mask>', '<cls>', '<sep>', '<eos>']



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



def        SCREAMING_SNAKE_CASE_  (	snake_case_				:	Tuple   ,     snake_case_				:	Any							)       ->	Union[str, Any]:
 SCREAMING_SNAKE_CASE  :  List[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        SCREAMING_SNAKE_CASE_  (	snake_case_				:	str   ,     snake_case_				:	str   ,     snake_case_				:	bool   ,     snake_case_				:	str   ,     snake_case_				:	str							)       ->	int:
 if model.startswith('esmfold'							):
  SCREAMING_SNAKE_CASE  :  int								=		MODEL_MAPPING[model]()
 else:
  SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE  :  List[Any]								=		MODEL_MAPPING[model]()
 esm.eval()  # disable dropout

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

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

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

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

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

 # Now let's copy all the weights.
 # Embeddings
 SCREAMING_SNAKE_CASE  :  Dict								=		original_esm_model.embed_tokens.weight
 if position_embedding_type == "absolute":
  SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		original_esm_model.embed_positions.weight

 if config.emb_layer_norm_before:
  SCREAMING_SNAKE_CASE  :  Optional[int]								=		original_esm_model.emb_layer_norm_before.weight
  SCREAMING_SNAKE_CASE  :  Any								=		original_esm_model.emb_layer_norm_before.bias

 SCREAMING_SNAKE_CASE  :  Any								=		original_esm_model.emb_layer_norm_after.weight
 SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		original_esm_model.emb_layer_norm_after.bias

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

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

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

  if getattr(esm_layer.self_attn   ,     'rot_emb'   ,     snake_case_							) 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.
   SCREAMING_SNAKE_CASE  :  int								=		esm_layer.self_attn.rot_emb.inv_freq

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

  # self-attention output
  SCREAMING_SNAKE_CASE  :  EsmSelfOutput								=		layer.attention.output
  assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
  SCREAMING_SNAKE_CASE  :  List[str]								=		esm_layer.self_attn.out_proj.weight
  SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		esm_layer.self_attn.out_proj.bias

  # intermediate
  SCREAMING_SNAKE_CASE  :  EsmIntermediate								=		layer.intermediate
  assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
  SCREAMING_SNAKE_CASE  :  List[str]								=		esm_layer.fca.weight
  SCREAMING_SNAKE_CASE  :  List[str]								=		esm_layer.fca.bias

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

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

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

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

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

  del esm  # Free up some memory before continuing

 print(f"""Saving tokenizer to {pytorch_dump_folder_path}"""							)
 hf_tokenizer.save_pretrained(snake_case_							)

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


if __name__ == "__main__":
   __UpperCAmelCase												=   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												=   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'''

def        SCREAMING_SNAKE_CASE_  (	snake_case_				:	int   ,     snake_case_				:	list[tuple[int, int, int]]							)       ->	list[tuple[int, int, int]]:
 SCREAMING_SNAKE_CASE  :  List[Any]								=		sorted(snake_case_   ,     key=lambda snake_case_							: edge[2]							)

 SCREAMING_SNAKE_CASE  :  str								=		list(range(snake_case_							)							)

 def find_parent(snake_case_				:	Union[str, Any]							):
  if i != parent[i]:
   SCREAMING_SNAKE_CASE  :  int								=		find_parent(parent[i]							)
  return parent[i]

 SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		0
 SCREAMING_SNAKE_CASE  :  Any								=		[]

 for edge in edges:
  SCREAMING_SNAKE_CASE  :  Dict								=		find_parent(edge[0]							)
  SCREAMING_SNAKE_CASE  :  Optional[int]								=		find_parent(edge[1]							)
  if parent_a != parent_b:
   minimum_spanning_tree_cost += edge[2]
   minimum_spanning_tree.append(snake_case_							)
   SCREAMING_SNAKE_CASE  :  Tuple								=		parent_b

 return minimum_spanning_tree


if __name__ == "__main__":  # pragma: no cover
   __UpperCAmelCase						,						__UpperCAmelCase												=   list(map(int, input().strip().split()))
   __UpperCAmelCase												=   []

   for _ in range(num_edges):
      __UpperCAmelCase						,						__UpperCAmelCase						,						__UpperCAmelCase												=   (int(x) for x in input().strip().split())
      edges.append((nodea, nodea, cost))

   kruskal(num_nodes, edges)


