


def   lowerCAmelCase_ (				__lowerCAmelCase		,					__lowerCAmelCase						)->  list[tuple[int, int, int]]:






      '''simple docstring'''





      UpperCAmelCase    : List[Any]      =sorted(__lowerCAmelCase		,					key=lambda __lowerCAmelCase						: edge[2]						)

      UpperCAmelCase    : Union[str, Any]      =list(range(__lowerCAmelCase						)						)

      def find_parent(__lowerCAmelCase						):
            if i != parent[i]:
                  UpperCAmelCase    : Optional[int]      =find_parent(parent[i]						)
            return parent[i]

      UpperCAmelCase    : Optional[Any]      =0
      UpperCAmelCase    : List[str]      =[]

      for edge in edges:
            UpperCAmelCase    : Dict      =find_parent(edge[0]						)
            UpperCAmelCase    : Optional[Any]      =find_parent(edge[1]						)
            if parent_a != parent_b:
                  minimum_spanning_tree_cost += edge[2]
                  minimum_spanning_tree.append(__lowerCAmelCase						)
                  UpperCAmelCase    : Optional[Any]      =parent_b

      return minimum_spanning_tree


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

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

     kruskal(num_nodes, edges)

import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


__snake_case			      =    logging.getLogger(__name__)

require_version('''pytorch_lightning>=1.0.4''')

__snake_case			      =    {
    '''base''': AutoModel,
    '''sequence-classification''': AutoModelForSequenceClassification,
    '''question-answering''': AutoModelForQuestionAnswering,
    '''pretraining''': AutoModelForPreTraining,
    '''token-classification''': AutoModelForTokenClassification,
    '''language-modeling''': AutoModelWithLMHead,
    '''summarization''': AutoModelForSeqaSeqLM,
    '''translation''': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
__snake_case			      =    {
    '''linear''': get_linear_schedule_with_warmup,
    '''cosine''': get_cosine_schedule_with_warmup,
    '''cosine_w_restarts''': get_cosine_with_hard_restarts_schedule_with_warmup,
    '''polynomial''': get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
__snake_case			      =    sorted(arg_to_scheduler.keys())
__snake_case			      =    '''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''







class 			__snake_case   (     pl.LightningModule   ):
 def __init__(  self		,   snake_case__		,   snake_case__=None		,   snake_case__="base"		,   snake_case__=None		,   snake_case__=None		,   snake_case__=None		,   **snake_case__		,   )	->							List[str]:
       '''simple docstring'''





       super().__init__()
       # TODO: move to self.save_hyperparameters()
       # self.save_hyperparameters()
       # can also expand arguments into trainer signature for easier reading

       self.save_hyperparameters(snake_case__						)
       UpperCAmelCase    : Union[str, Any]      =0
       UpperCAmelCase    : List[Any]      =Path(self.hparams.output_dir						)
       UpperCAmelCase    : List[str]      =self.hparams.cache_dir if self.hparams.cache_dir else None
       if config is None:
             UpperCAmelCase    : Union[str, Any]      =AutoConfig.from_pretrained(
                 self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path		,   **({'''num_labels''': num_labels} if num_labels is not None else {})		,   cache_dir=snake_case__		,   **snake_case__		,   )
       else:
             UpperCAmelCase    : PretrainedConfig      =config

       UpperCAmelCase    : List[str]      =('''encoder_layerdrop''', '''decoder_layerdrop''', '''dropout''', '''attention_dropout''')
       for p in extra_model_params:
             if getattr(self.hparams		,   snake_case__		,   snake_case__						):
                   assert hasattr(self.config		,   snake_case__						), f'''model config doesn\'t have a `{p}` attribute'''
                   setattr(self.config		,   snake_case__		,   getattr(self.hparams		,   snake_case__						)						)

       if tokenizer is None:
             UpperCAmelCase    : Dict      =AutoTokenizer.from_pretrained(
                 self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path		,   cache_dir=snake_case__		,   )
       else:
             UpperCAmelCase    : PreTrainedTokenizer      =tokenizer
       UpperCAmelCase    : Dict      =MODEL_MODES[mode]
       if model is None:
             UpperCAmelCase    : str      =self.model_type.from_pretrained(
                 self.hparams.model_name_or_path		,   from_tf=bool('''.ckpt''' in self.hparams.model_name_or_path						)		,   config=self.config		,   cache_dir=snake_case__		,   )
       else:
             UpperCAmelCase    : Union[str, Any]      =model
 def 	UpperCAmelCase__       (  self		,   *snake_case__		,   **snake_case__						)	->							str:
       '''simple docstring'''





       UpperCAmelCase    : Optional[int]      =self.model_type.from_pretrained(*snake_case__		,   **snake_case__						)
 def 	UpperCAmelCase__       (  self						)	->							Optional[Any]:
       '''simple docstring'''





       UpperCAmelCase    : Optional[Any]      =arg_to_scheduler[self.hparams.lr_scheduler]
       UpperCAmelCase    : Dict      =get_schedule_func(
           self.opt		,   num_warmup_steps=self.hparams.warmup_steps		,   num_training_steps=self.total_steps()						)
       UpperCAmelCase    : int      ={'''scheduler''': scheduler, '''interval''': '''step''', '''frequency''': 1}
       return scheduler
 def 	UpperCAmelCase__       (  self						)	->							Tuple:
       '''simple docstring'''





       UpperCAmelCase    : List[str]      =self.model
       UpperCAmelCase    : Dict      =['''bias''', '''LayerNorm.weight''']
       UpperCAmelCase    : Dict      =[
           {
               '''params''': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay						)],
               '''weight_decay''': self.hparams.weight_decay,
           },
           {
               '''params''': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay						)],
               '''weight_decay''': 0.0,
           },
       ]
       if self.hparams.adafactor:
             UpperCAmelCase    : Any      =Adafactor(
                 snake_case__		,   lr=self.hparams.learning_rate		,   scale_parameter=snake_case__		,   relative_step=snake_case__						)

       else:
             UpperCAmelCase    : Optional[int]      =AdamW(
                 snake_case__		,   lr=self.hparams.learning_rate		,   eps=self.hparams.adam_epsilon						)
       UpperCAmelCase    : int      =optimizer

       UpperCAmelCase    : Union[str, Any]      =self.get_lr_scheduler()

       return [optimizer], [scheduler]
 def 	UpperCAmelCase__       (  self		,   snake_case__		,   snake_case__						)	->							Tuple:
       '''simple docstring'''





       return self.validation_step(snake_case__		,   snake_case__						)
 def 	UpperCAmelCase__       (  self		,   snake_case__						)	->							Optional[int]:
       '''simple docstring'''





       return self.validation_end(snake_case__						)
 def 	UpperCAmelCase__       (  self						)	->							int:
       '''simple docstring'''





       UpperCAmelCase    : str      =max(1		,   self.hparams.gpus						)  # TODO: consider num_tpu_cores
       UpperCAmelCase    : Dict      =self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
       return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs
 def 	UpperCAmelCase__       (  self		,   snake_case__						)	->							Union[str, Any]:
       '''simple docstring'''





       if mode == "test":
             UpperCAmelCase    : str      =len(self.test_dataloader().dataset						)
       else:
             UpperCAmelCase    : Union[str, Any]      =self.get_dataloader('''train'''		,   self.hparams.train_batch_size		,   shuffle=snake_case__						)
             UpperCAmelCase    : str      =len(self.train_dataloader().dataset						)
 def 	UpperCAmelCase__       (  self		,   snake_case__		,   snake_case__		,   snake_case__ = False						)	->							List[str]:
       '''simple docstring'''





       raise NotImplementedError('''You must implement this for your task'''						)
 def 	UpperCAmelCase__       (  self						)	->							Union[str, Any]:
       '''simple docstring'''





       return self.train_loader
 def 	UpperCAmelCase__       (  self						)	->							List[Any]:
       '''simple docstring'''





       return self.get_dataloader('''dev'''		,   self.hparams.eval_batch_size		,   shuffle=snake_case__						)
 def 	UpperCAmelCase__       (  self						)	->							Tuple:
       '''simple docstring'''





       return self.get_dataloader('''test'''		,   self.hparams.eval_batch_size		,   shuffle=snake_case__						)
 def 	UpperCAmelCase__       (  self		,   snake_case__						)	->							Tuple:
       '''simple docstring'''





       return os.path.join(
           self.hparams.data_dir		,   '''cached_{}_{}_{}'''.format(
               snake_case__		,   list(filter(snake_case__		,   self.hparams.model_name_or_path.split('''/'''						)						)						).pop()		,   str(self.hparams.max_seq_length						)		,   )		,   )
 @pl.utilities.rank_zero_only
 def 	UpperCAmelCase__       (  self		,   snake_case__						)	->							None:
       '''simple docstring'''





       UpperCAmelCase    : List[str]      =self.output_dir.joinpath('''best_tfmr'''						)
       UpperCAmelCase    : int      =self.step_count
       self.model.save_pretrained(snake_case__						)
       self.tokenizer.save_pretrained(snake_case__						)






 @staticmethod
 def 	UpperCAmelCase__       (  snake_case__		,   snake_case__						)	->							Any:
       '''simple docstring'''





       parser.add_argument(
           '''--model_name_or_path'''		,   default=snake_case__		,   type=snake_case__		,   required=snake_case__		,   help='''Path to pretrained model or model identifier from huggingface.co/models'''		,   )
       parser.add_argument(
           '''--config_name'''		,   default=''''''		,   type=snake_case__		,   help='''Pretrained config name or path if not the same as model_name'''						)
       parser.add_argument(
           '''--tokenizer_name'''		,   default=snake_case__		,   type=snake_case__		,   help='''Pretrained tokenizer name or path if not the same as model_name'''		,   )
       parser.add_argument(
           '''--cache_dir'''		,   default=''''''		,   type=snake_case__		,   help='''Where do you want to store the pre-trained models downloaded from huggingface.co'''		,   )
       parser.add_argument(
           '''--encoder_layerdrop'''		,   type=snake_case__		,   help='''Encoder layer dropout probability (Optional). Goes into model.config'''		,   )
       parser.add_argument(
           '''--decoder_layerdrop'''		,   type=snake_case__		,   help='''Decoder layer dropout probability (Optional). Goes into model.config'''		,   )
       parser.add_argument(
           '''--dropout'''		,   type=snake_case__		,   help='''Dropout probability (Optional). Goes into model.config'''		,   )
       parser.add_argument(
           '''--attention_dropout'''		,   type=snake_case__		,   help='''Attention dropout probability (Optional). Goes into model.config'''		,   )
       parser.add_argument('''--learning_rate'''		,   default=5e-5		,   type=snake_case__		,   help='''The initial learning rate for Adam.'''						)
       parser.add_argument(
           '''--lr_scheduler'''		,   default='''linear'''		,   choices=snake_case__		,   metavar=snake_case__		,   type=snake_case__		,   help='''Learning rate scheduler'''		,   )
       parser.add_argument('''--weight_decay'''		,   default=0.0		,   type=snake_case__		,   help='''Weight decay if we apply some.'''						)
       parser.add_argument('''--adam_epsilon'''		,   default=1e-8		,   type=snake_case__		,   help='''Epsilon for Adam optimizer.'''						)
       parser.add_argument('''--warmup_steps'''		,   default=0		,   type=snake_case__		,   help='''Linear warmup over warmup_steps.'''						)
       parser.add_argument('''--num_workers'''		,   default=4		,   type=snake_case__		,   help='''kwarg passed to DataLoader'''						)
       parser.add_argument('''--num_train_epochs'''		,   dest='''max_epochs'''		,   default=3		,   type=snake_case__						)
       parser.add_argument('''--train_batch_size'''		,   default=32		,   type=snake_case__						)
       parser.add_argument('''--eval_batch_size'''		,   default=32		,   type=snake_case__						)
       parser.add_argument('''--adafactor'''		,   action='''store_true'''						)







class 			__snake_case   (     pl.Callback   ):
 def 	UpperCAmelCase__       (  self		,   snake_case__		,   snake_case__						)	->							List[Any]:
       '''simple docstring'''





       UpperCAmelCase    : Tuple      =trainer.lr_schedulers[0]['''scheduler''']
       UpperCAmelCase    : Dict      ={f'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr()						)}
       pl_module.logger.log_metrics(snake_case__						)
 def 	UpperCAmelCase__       (  self		,   snake_case__		,   snake_case__						)	->							Optional[int]:
       '''simple docstring'''





       rank_zero_info('''***** Validation results *****'''						)
       UpperCAmelCase    : Union[str, Any]      =trainer.callback_metrics
       # Log results
       for key in sorted(snake_case__						):
             if key not in ["log", "progress_bar"]:
                   rank_zero_info('''{} = {}\n'''.format(snake_case__		,   str(metrics[key]						)						)						)






 def 	UpperCAmelCase__       (  self		,   snake_case__		,   snake_case__						)	->							Tuple:
       '''simple docstring'''





       rank_zero_info('''***** Test results *****'''						)
       UpperCAmelCase    : List[str]      =trainer.callback_metrics
       # Log and save results to file
       UpperCAmelCase    : List[Any]      =os.path.join(pl_module.hparams.output_dir		,   '''test_results.txt'''						)
       with open(snake_case__		,   '''w'''						) as writer:
             for key in sorted(snake_case__						):
                   if key not in ["log", "progress_bar"]:
                         rank_zero_info('''{} = {}\n'''.format(snake_case__		,   str(metrics[key]						)						)						)
                         writer.write('''{} = {}\n'''.format(snake_case__		,   str(metrics[key]						)						)						)





def   lowerCAmelCase_ (				__lowerCAmelCase		,					__lowerCAmelCase						)->  None:






      '''simple docstring'''





      parser.add_argument(
          '''--output_dir'''		,					default=__lowerCAmelCase		,					type=__lowerCAmelCase		,					required=__lowerCAmelCase		,					help='''The output directory where the model predictions and checkpoints will be written.'''		,					)
      parser.add_argument(
          '''--fp16'''		,					action='''store_true'''		,					help='''Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit'''		,					)

      parser.add_argument(
          '''--fp16_opt_level'''		,					type=__lowerCAmelCase		,					default='''O2'''		,					help=(
              '''For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\'].'''
              '''See details at https://nvidia.github.io/apex/amp.html'''
          )		,					)
      parser.add_argument('''--n_tpu_cores'''		,					dest='''tpu_cores'''		,					type=__lowerCAmelCase						)
      parser.add_argument('''--max_grad_norm'''		,					dest='''gradient_clip_val'''		,					default=1.0		,					type=__lowerCAmelCase		,					help='''Max gradient norm'''						)
      parser.add_argument('''--do_train'''		,					action='''store_true'''		,					help='''Whether to run training.'''						)
      parser.add_argument('''--do_predict'''		,					action='''store_true'''		,					help='''Whether to run predictions on the test set.'''						)
      parser.add_argument(
          '''--gradient_accumulation_steps'''		,					dest='''accumulate_grad_batches'''		,					type=__lowerCAmelCase		,					default=1		,					help='''Number of updates steps to accumulate before performing a backward/update pass.'''		,					)
      parser.add_argument('''--seed'''		,					type=__lowerCAmelCase		,					default=42		,					help='''random seed for initialization'''						)
      parser.add_argument(
          '''--data_dir'''		,					default=__lowerCAmelCase		,					type=__lowerCAmelCase		,					required=__lowerCAmelCase		,					help='''The input data dir. Should contain the training files for the CoNLL-2003 NER task.'''		,					)





def   lowerCAmelCase_ (				__lowerCAmelCase		,					__lowerCAmelCase		,					__lowerCAmelCase=None		,					__lowerCAmelCase=True		,					__lowerCAmelCase=[]		,					__lowerCAmelCase=None		,					__lowerCAmelCase=None		,					**__lowerCAmelCase		,					)->  Optional[int]:






      '''simple docstring'''





      pl.seed_everything(args.seed						)

      # init model
      UpperCAmelCase    : List[Any]      =Path(model.hparams.output_dir						)
      odir.mkdir(exist_ok=__lowerCAmelCase						)

      # add custom checkpoints
      if checkpoint_callback is None:
            UpperCAmelCase    : Any      =pl.callbacks.ModelCheckpoint(
                filepath=args.output_dir		,					prefix='''checkpoint'''		,					monitor='''val_loss'''		,					mode='''min'''		,					save_top_k=1						)
      if early_stopping_callback:
            extra_callbacks.append(__lowerCAmelCase						)
      if logging_callback is None:
            UpperCAmelCase    : int      =LoggingCallback()

      UpperCAmelCase    : Any      ={}

      # TODO: remove with PyTorch 1.6 since pl uses native amp
      if args.fpaa:
            UpperCAmelCase    : List[str]      =16
            UpperCAmelCase    : List[Any]      =args.fpaa_opt_level

      if args.gpus > 1:
            UpperCAmelCase    : str      ='''ddp'''

      UpperCAmelCase    : str      =args.accumulate_grad_batches
      UpperCAmelCase    : List[Any]      =extra_train_kwargs.get('''accelerator'''		,					__lowerCAmelCase						)
      UpperCAmelCase    : Optional[int]      =extra_train_kwargs.get('''profiler'''		,					__lowerCAmelCase						)

      UpperCAmelCase    : str      =pl.Trainer.from_argparse_args(
          __lowerCAmelCase		,					weights_summary=__lowerCAmelCase		,					callbacks=[logging_callback] + extra_callbacks		,					logger=__lowerCAmelCase		,					checkpoint_callback=__lowerCAmelCase		,					**__lowerCAmelCase		,					)

      if args.do_train:
            trainer.fit(__lowerCAmelCase						)

      return trainer
