





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_rag import save_json




def  lowerCAmelCase_ (							UpperCamelCase_  )     ->		int:
							UpperCamelCase_          =  filter(lambda UpperCamelCase_  : p.requires_grad					,		model.parameters()  )
							UpperCamelCase_          =  sum([np.prod(p.size()  ) for p in model_parameters]  )
							return params


_UpperCAmelCase        =				logging.getLogger(__name__)




def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Optional[int]:



							if metric == "rouge2":
														UpperCamelCase_          =  '{val_avg_rouge2:.4f}-{step_count}'
							elif metric == "bleu":
														UpperCamelCase_          =  '{val_avg_bleu:.4f}-{step_count}'
							elif metric == "em":
														UpperCamelCase_          =  '{val_avg_em:.4f}-{step_count}'
							else:
														raise NotImplementedError(
														    F'''seq2seq callbacks only support rouge2 and bleu, got {metric}, You can make your own by adding to this'''
														    " function."  )

							UpperCamelCase_          =  ModelCheckpoint(
							    dirpath=_lowerCamelCase					,		filename=_lowerCamelCase					,		monitor=F'''val_{metric}'''					,		mode="max"					,		save_top_k=3					,		every_n_epochs=1					,		)
							return checkpoint_callback




def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Union[str, Any]:
							return EarlyStopping(
							    monitor=F'''val_{metric}'''					,		mode="min" if "loss" in metric else "max"					,		patience=_lowerCamelCase					,		verbose=_lowerCamelCase					,		)







class 							_UpperCamelCase					(	pl.Callback     ):

							def lowercase (		self: str					,   _SCREAMING_SNAKE_CASE: int					,   _SCREAMING_SNAKE_CASE: Optional[Any]		)       -> Any:




														"""simple docstring"""


														UpperCamelCase_          =  {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 lowercase (		self: Optional[int]					,   _SCREAMING_SNAKE_CASE: int					,   _SCREAMING_SNAKE_CASE: List[str]					,   _SCREAMING_SNAKE_CASE: Dict					,   _SCREAMING_SNAKE_CASE: Dict=True		)       -> List[str]:




														"""simple docstring"""


														logger.info(f'''***** {type_path} results at step {trainer.global_step:05d} *****'''		)
														UpperCamelCase_          =  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
														UpperCamelCase_          =  Path(pl_module.hparams.output_dir		)
														if type_path == "test":
																					UpperCamelCase_          =  od / 'test_results.txt'
																					UpperCamelCase_          =  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.
																					UpperCamelCase_          =  od / f'''{type_path}_results/{trainer.global_step:05d}.txt'''
																					UpperCamelCase_          =  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
																												UpperCamelCase_          =  metrics[key]
																												if isinstance(_A					,   torch.Tensor		):
																																			UpperCamelCase_          =  val.item()
																												UpperCamelCase_          =  f'''{key}: {val:.6f}\n'''
																												writer.write(_A		)

														if not save_generations:
																					return

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

							@rank_zero_only
							def lowercase (		self: Optional[int]					,   _SCREAMING_SNAKE_CASE: List[Any]					,   _SCREAMING_SNAKE_CASE: Tuple		)       -> Tuple:




														"""simple docstring"""


														try:
																					UpperCamelCase_          =  pl_module.model.model.num_parameters()
														except AttributeError:
																					UpperCamelCase_          =  pl_module.model.num_parameters()

														UpperCamelCase_          =  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 lowercase (		self: List[str]					,   _SCREAMING_SNAKE_CASE: Optional[int]					,   _SCREAMING_SNAKE_CASE: List[str]		)       -> str:




														"""simple docstring"""


														save_json(pl_module.metrics					,   pl_module.metrics_save_path		)
														return self._write_logs(_A					,   _A					,   "test"		)



							@rank_zero_only
							def lowercase (		self: List[Any]					,   _SCREAMING_SNAKE_CASE: Any					,   _SCREAMING_SNAKE_CASE: List[str]		)       -> Dict:




														"""simple docstring"""


														save_json(pl_module.metrics					,   pl_module.metrics_save_path		)
														# Uncommenting this will save val generations
														# return self._write_logs(trainer, pl_module, "valid")








import re
import string

import numpy as np

import datasets


_UpperCAmelCase        =				'\nReturns the rate at which the input predicted strings exactly match their references, ignoring any strings input as part of the regexes_to_ignore list.\n'

_UpperCAmelCase        =				'\nArgs:\n    predictions: List of predicted texts.\n    references: List of reference texts.\n    regexes_to_ignore: List, defaults to None. Regex expressions of characters to\n        ignore when calculating the exact matches. Note: these regexes are removed\n        from the input data before the changes based on the options below (e.g. ignore_case,\n        ignore_punctuation, ignore_numbers) are applied.\n    ignore_case: Boolean, defaults to False. If true, turns everything\n        to lowercase so that capitalization differences are ignored.\n    ignore_punctuation: Boolean, defaults to False. If true, removes all punctuation before\n        comparing predictions and references.\n    ignore_numbers: Boolean, defaults to False. If true, removes all punctuation before\n        comparing predictions and references.\nReturns:\n    exact_match: Dictionary containing exact_match rate. Possible values are between 0.0 and 100.0, inclusive.\nExamples:\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds)\n    >>> print(round(results["exact_match"], 1))\n    25.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell"], ignore_case=True, ignore_punctuation=True)\n    >>> print(round(results["exact_match"], 1))\n    50.0\n\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell", "YELL"], ignore_case=True, ignore_punctuation=True)\n    >>> print(round(results["exact_match"], 1))\n    75.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell", "YELL"], ignore_case=True, ignore_punctuation=True, ignore_numbers=True)\n    >>> print(round(results["exact_match"], 1))\n    100.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["The cat sat on the mat.", "Theaters are great.", "It\'s like comparing oranges and apples."]\n    >>> preds = ["The cat sat on the mat?", "Theaters are great.", "It\'s like comparing apples and oranges."]\n    >>> results = exact_match.compute(references=refs, predictions=preds)\n    >>> print(round(results["exact_match"], 1))\n    33.3\n\n'

_UpperCAmelCase        =				'\n'



@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION  ,						_KWARGS_DESCRIPTION     )
class 							_UpperCamelCase					(	datasets.Metric     ):

       def lowercase (		self: Any		)       -> List[str]:




              """simple docstring"""


              return datasets.MetricInfo(
                  description=_DESCRIPTION					,   citation=_CITATION					,   inputs_description=_KWARGS_DESCRIPTION					,   features=datasets.Features(
                      {
                          "predictions": datasets.Value("string"					,   id="sequence"		),
                          "references": datasets.Value("string"					,   id="sequence"		),
                      }		)					,   reference_urls=[]					,   )



       def lowercase (		self: Any					,   _SCREAMING_SNAKE_CASE: Any					,   _SCREAMING_SNAKE_CASE: Tuple					,   _SCREAMING_SNAKE_CASE: List[Any]=None					,   _SCREAMING_SNAKE_CASE: Tuple=False					,   _SCREAMING_SNAKE_CASE: Optional[int]=False					,   _SCREAMING_SNAKE_CASE: Any=False					,   )       -> Optional[Any]:




              """simple docstring"""


              if regexes_to_ignore is not None:
                     for s in regexes_to_ignore:
                            UpperCamelCase_          =  np.array([re.sub(_SCREAMING_SNAKE_CASE					,   ""					,   _SCREAMING_SNAKE_CASE		) for x in predictions]		)
                            UpperCamelCase_          =  np.array([re.sub(_SCREAMING_SNAKE_CASE					,   ""					,   _SCREAMING_SNAKE_CASE		) for x in references]		)
              else:
                     UpperCamelCase_          =  np.asarray(_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  np.asarray(_SCREAMING_SNAKE_CASE		)

              if ignore_case:
                     UpperCamelCase_          =  np.char.lower(_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  np.char.lower(_SCREAMING_SNAKE_CASE		)

              if ignore_punctuation:
                     UpperCamelCase_          =  string.punctuation.maketrans(""					,   ""					,   string.punctuation		)
                     UpperCamelCase_          =  np.char.translate(_SCREAMING_SNAKE_CASE					,   table=_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  np.char.translate(_SCREAMING_SNAKE_CASE					,   table=_SCREAMING_SNAKE_CASE		)

              if ignore_numbers:
                     UpperCamelCase_          =  string.digits.maketrans(""					,   ""					,   string.digits		)
                     UpperCamelCase_          =  np.char.translate(_SCREAMING_SNAKE_CASE					,   table=_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  np.char.translate(_SCREAMING_SNAKE_CASE					,   table=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  predictions == references

              return {"exact_match": np.mean(_SCREAMING_SNAKE_CASE		) * 100}

