




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 					lowerCamelCase__							(			__lowercase							):
     snake_case :			List[str]					   = filter(lambda __lowercase							: p.requires_grad    ,      model.parameters()							)
     snake_case :			List[Any]					   = sum([np.prod(p.size()							) for p in model_parameters]							)
     return params


lowercase			:      Union[str, Any]			=					logging.getLogger(__name__)







class 				_a   (pl.Callback  ):





     '''simple docstring'''



     def  snake_case_			(							self       ,__a       ,__a       )			->							str:
          snake_case :			int					   = {F'''lr_group_{i}''': param["""lr"""] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups       )}
          pl_module.logger.log_metrics(__a       )







     @rank_zero_only
     def  snake_case_			(							self       ,__a       ,__a       ,__a       ,__a=True       )			->							None:
          logger.info(F'''***** {type_path} results at step {trainer.global_step:05d} *****'''       )
          snake_case :			Tuple					   = 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
          snake_case :			Dict					   = Path(pl_module.hparams.output_dir       )
          if type_path == "test":
               snake_case :			Tuple					   = od / """test_results.txt"""
               snake_case :			Dict					   = 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.
               snake_case :			int					   = od / F'''{type_path}_results/{trainer.global_step:05d}.txt'''
               snake_case :			List[str]					   = od / F'''{type_path}_generations/{trainer.global_step:05d}.txt'''
               results_file.parent.mkdir(exist_ok=__a       )
               generations_file.parent.mkdir(exist_ok=__a       )
          with open(__a       ,"""a+"""       ) as writer:
               for key in sorted(__a       ):
                    if key in ["log", "progress_bar", "preds"]:
                         continue
                    snake_case :			str					   = metrics[key]
                    if isinstance(__a       ,torch.Tensor       ):
                         snake_case :			Union[str, Any]					   = val.item()
                    snake_case :			int					   = F'''{key}: {val:.6f}\n'''
                    writer.write(__a       )

          if not save_generations:
               return

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







     @rank_zero_only
     def  snake_case_			(							self       ,__a       ,__a       )			->							Union[str, Any]:
          try:
               snake_case :			int					   = pl_module.model.model.num_parameters()
          except AttributeError:
               snake_case :			Dict					   = pl_module.model.num_parameters()

          snake_case :			str					   = count_trainable_parameters(__a       )
          # 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  snake_case_			(							self       ,__a       ,__a       )			->							Dict:
          save_json(pl_module.metrics       ,pl_module.metrics_save_path       )
          return self._write_logs(__a       ,__a       ,"""test"""       )







     @rank_zero_only
     def  snake_case_			(							self       ,__a       ,__a       )			->							int:
          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 					lowerCamelCase__							(			__lowercase    ,      __lowercase    ,      __lowercase=1    ,      __lowercase=False							):
     if metric == "rouge2":
          snake_case :			Union[str, Any]					   = """{val_avg_rouge2:.4f}-{step_count}"""
     elif metric == "bleu":
          snake_case :			List[Any]					   = """{val_avg_bleu:.4f}-{step_count}"""
     elif metric == "loss":
          snake_case :			Dict					   = """{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."""							)

     snake_case :			Dict					   = ModelCheckpoint(
         dirpath=__lowercase    ,      filename=__lowercase    ,      monitor=F'''val_{metric}'''    ,      mode="""min""" if """loss""" in metric else """max"""    ,      save_top_k=__lowercase    ,      )
     return checkpoint_callback







def 					lowerCamelCase__							(			__lowercase    ,      __lowercase							):
     return EarlyStopping(
         monitor=F'''val_{metric}'''    ,      mode="""min""" if """loss""" in metric else """max"""    ,      patience=__lowercase    ,      verbose=__lowercase    ,      )









'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowercase			:      Any			=					{"""configuration_visual_bert""": ["""VISUAL_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """VisualBertConfig"""]}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      lowercase			:      Optional[int]			=					[
          """VISUAL_BERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """VisualBertForMultipleChoice""",
          """VisualBertForPreTraining""",
          """VisualBertForQuestionAnswering""",
          """VisualBertForRegionToPhraseAlignment""",
          """VisualBertForVisualReasoning""",
          """VisualBertLayer""",
          """VisualBertModel""",
          """VisualBertPreTrainedModel""",
      ]


if TYPE_CHECKING:
      from .configuration_visual_bert import VISUAL_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, VisualBertConfig

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_visual_bert import (
                VISUAL_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                VisualBertForMultipleChoice,
                VisualBertForPreTraining,
                VisualBertForQuestionAnswering,
                VisualBertForRegionToPhraseAlignment,
                VisualBertForVisualReasoning,
                VisualBertLayer,
                VisualBertModel,
                VisualBertPreTrainedModel,
            )


else:
      import sys

      lowercase			:      List[Any]			=					_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



