


"""simple docstring"""





import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only

from utils import save_json


def   _snake_case    (    _snake_case					:      List[str]      ):
      lowerCAmelCase       :						Dict							= filter(lambda _snake_case      : p.requires_grad			, model.parameters()      )
      lowerCAmelCase       :						Any							= sum([np.prod(p.size()      ) for p in model_parameters]      )
      return params


snake_case__		:  List[str]     					=					logging.getLogger(__name__)



class 		snake_case_(							pl.Callback      ):






     def   lowerCamelCase__  (						self	:       List[str]       ,					UpperCamelCase_	:       int       ,					UpperCamelCase_	:       Tuple					):
           lowerCAmelCase       :						int							= {F'''lr_group_{i}''': param['''lr'''] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups					)}
           pl_module.logger.log_metrics(UpperCamelCase_					)






     @rank_zero_only
     def   lowerCamelCase__  (						self	:       Union[str, Any]       ,					UpperCamelCase_	:       pl.Trainer       ,					UpperCamelCase_	:       pl.LightningModule       ,					UpperCamelCase_	:       str       ,					UpperCamelCase_	:       List[Any]=True					):
           logger.info(F'''***** {type_path} results at step {trainer.global_step:05d} *****'''					)
           lowerCAmelCase       :						int							= trainer.callback_metrics
           trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ['''log''', '''progress_bar''', '''preds''']}					)
           # Log results
           lowerCAmelCase       :						Union[str, Any]							= Path(pl_module.hparams.output_dir					)
           if type_path == "test":
                 lowerCAmelCase       :						Tuple							= od / '''test_results.txt'''
                 lowerCAmelCase       :						Optional[Any]							= od / '''test_generations.txt'''
           else:
                 # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
                 # If people want this it will be easy enough to add back.
                 lowerCAmelCase       :						str							= od / F'''{type_path}_results/{trainer.global_step:05d}.txt'''
                 lowerCAmelCase       :						List[Any]							= od / F'''{type_path}_generations/{trainer.global_step:05d}.txt'''
                 results_file.parent.mkdir(exist_ok=UpperCamelCase_					)
                 generations_file.parent.mkdir(exist_ok=UpperCamelCase_					)
           with open(UpperCamelCase_       ,					'''a+'''					) as writer:
                 for key in sorted(UpperCamelCase_					):
                       if key in ["log", "progress_bar", "preds"]:
                             continue
                       lowerCAmelCase       :						int							= metrics[key]
                       if isinstance(UpperCamelCase_       ,					torch.Tensor					):
                             lowerCAmelCase       :						Dict							= val.item()
                       lowerCAmelCase       :						List[str]							= F'''{key}: {val:.6f}\n'''
                       writer.write(UpperCamelCase_					)

           if not save_generations:
                 return

           if "preds" in metrics:
                 lowerCAmelCase       :						Dict							= '''\n'''.join(metrics['''preds''']					)
                 generations_file.open('''w+'''					).write(UpperCamelCase_					)






     @rank_zero_only
     def   lowerCamelCase__  (						self	:       Tuple       ,					UpperCamelCase_	:       List[Any]       ,					UpperCamelCase_	:       Dict					):
           try:
                 lowerCAmelCase       :						Union[str, Any]							= pl_module.model.model.num_parameters()
           except AttributeError:
                 lowerCAmelCase       :						List[str]							= pl_module.model.num_parameters()

           lowerCAmelCase       :						List[str]							= count_trainable_parameters(UpperCamelCase_					)
           # mp stands for million parameters
           trainer.logger.log_metrics({'''n_params''': npars, '''mp''': npars / 1E6, '''grad_mp''': n_trainable_pars / 1E6}					)






     @rank_zero_only
     def   lowerCamelCase__  (						self	:       Dict       ,					UpperCamelCase_	:       pl.Trainer       ,					UpperCamelCase_	:       pl.LightningModule					):
           save_json(pl_module.metrics       ,					pl_module.metrics_save_path					)
           return self._write_logs(UpperCamelCase_       ,					UpperCamelCase_       ,					'''test'''					)






     @rank_zero_only
     def   lowerCamelCase__  (						self	:       Optional[Any]       ,					UpperCamelCase_	:       pl.Trainer       ,					UpperCamelCase_	:       Any					):
           save_json(pl_module.metrics       ,					pl_module.metrics_save_path					)
           # Uncommenting this will save val generations
           # return self._write_logs(trainer, pl_module, "valid")



def   _snake_case    (    _snake_case					:      Tuple			, _snake_case					:      str			, _snake_case					:      Union[str, Any]=1			, _snake_case					:      Optional[int]=False      ):
      if metric == "rouge2":
            lowerCAmelCase       :						int							= '''{val_avg_rouge2:.4f}-{step_count}'''
      elif metric == "bleu":
            lowerCAmelCase       :						Tuple							= '''{val_avg_bleu:.4f}-{step_count}'''
      elif metric == "loss":
            lowerCAmelCase       :						Tuple							= '''{val_avg_loss:.4f}-{step_count}'''
      else:
            raise NotImplementedError(
                f'''seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to'''
                ''' this function.'''      )

      lowerCAmelCase       :						Optional[Any]							= ModelCheckpoint(
          dirpath=_snake_case			, filename=_snake_case			, monitor=f'''val_{metric}'''			, mode='''min''' if '''loss''' in metric else '''max'''			, save_top_k=_snake_case			, )
      return checkpoint_callback






def   _snake_case    (    _snake_case					:      Tuple			, _snake_case					:      Dict      ):
      return EarlyStopping(
          monitor=f'''val_{metric}'''			, mode='''min''' if '''loss''' in metric else '''max'''			, patience=_snake_case			, verbose=_snake_case			, )






"""simple docstring"""





from doctest import testmod
from math import sqrt


def   _snake_case    (    _snake_case					:      int      ):
      lowerCAmelCase       :						list[int]							= []
      if num < 1:
            return facs
      facs.append(1      )
      if num == 1:
            return facs
      facs.append(_snake_case      )
      for i in range(2			, int(sqrt(_snake_case      )      ) + 1      ):
            if num % i == 0:  # If i is a factor of num
                  facs.append(_snake_case      )
                  lowerCAmelCase       :						Dict							= num // i  # num//i is the other factor of num
                  if d != i:  # If d and i are distinct
                        facs.append(_snake_case      )  # we have found another factor
      facs.sort()
      return facs


if __name__ == "__main__":
 testmod(name='''factors_of_a_number''', verbose=True)


