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



def       lowerCamelCase__		(			__snake_case	)							->		Dict:


					"""simple docstring"""





					_UpperCamelCase							 =  filter(lambda __snake_case	: p.requires_grad,     model.parameters()	)
					_UpperCamelCase							 =  sum([np.prod(p.size()	) for p in model_parameters]	)
					return params


_a														=		logging.getLogger(__name__)



def       lowerCamelCase__		(			__snake_case,     __snake_case	)							->		str:


					"""simple docstring"""





					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}'''
					elif metric == "loss":
										_UpperCamelCase							 =  '''{val_avg_loss:.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=__snake_case,     filename=__snake_case,     monitor=F'''val_{metric}''',     mode='''max''',     save_top_k=1,     every_n_epochs=1,     )

					return checkpoint_callback



def       lowerCamelCase__		(			__snake_case,     __snake_case	)							->		Dict:


					"""simple docstring"""





					return EarlyStopping(
					    monitor=F'''val_{metric}''',     mode='''min''' if '''loss''' in metric else '''max''',     patience=__snake_case,     verbose=__snake_case,     )






class     _UpperCAmelCase(    pl.Callback	):



		def 						UpperCAmelCase							(					self				,		__a				,		__a)       ->       List[str]:






							'''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 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a				,		__a=True)       ->       None:






							'''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 						UpperCAmelCase							(					self				,		__a				,		__a)       ->       List[str]:






							'''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 						UpperCAmelCase							(					self				,		__a				,		__a)       ->       Any:






							'''simple docstring'''
							save_json(pl_module.metrics				,		pl_module.metrics_save_path)
							return self._write_logs(__a				,		__a				,		'''test''')




		@rank_zero_only
		def 						UpperCAmelCase							(					self				,		__a				,		__a)       ->       Tuple:






							'''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")








"""simple docstring"""





from sklearn.metrics import matthews_corrcoef

import datasets


_a														=		"""
Compute the Matthews correlation coefficient (MCC)

The Matthews correlation coefficient is used in machine learning as a
measure of the quality of binary and multiclass classifications. It takes
into account true and false positives and negatives and is generally
regarded as a balanced measure which can be used even if the classes are of
very different sizes. The MCC is in essence a correlation coefficient value
between -1 and +1. A coefficient of +1 represents a perfect prediction, 0
an average random prediction and -1 an inverse prediction.  The statistic
is also known as the phi coefficient. [source: Wikipedia]
"""

_a														=		"""
Args:
    predictions (list of int): Predicted labels, as returned by a model.
    references (list of int): Ground truth labels.
    sample_weight (list of int, float, or bool): Sample weights. Defaults to `None`.
Returns:
    matthews_correlation (dict containing float): Matthews correlation.
Examples:
    Example 1, a basic example with only predictions and references as inputs:
        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")
        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],
        ...                                     predictions=[1, 2, 2, 0, 3, 3])
        >>> print(round(results['matthews_correlation'], 2))
        0.54

    Example 2, the same example as above, but also including sample weights:
        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")
        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],
        ...                                     predictions=[1, 2, 2, 0, 3, 3],
        ...                                     sample_weight=[0.5, 3, 1, 1, 1, 2])
        >>> print(round(results['matthews_correlation'], 2))
        0.1

    Example 3, the same example as above, but with sample weights that cause a negative correlation:
        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")
        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],
        ...                                     predictions=[1, 2, 2, 0, 3, 3],
        ...                                     sample_weight=[0.5, 1, 0, 0, 0, 1])
        >>> print(round(results['matthews_correlation'], 2))
        -0.25
"""

_a														=		"""\
@article{scikit-learn,
  title={Scikit-learn: Machine Learning in {P}ython},
  author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.
         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.
         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and
         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},
  journal={Journal of Machine Learning Research},
  volume={12},
  pages={2825--2830},
  year={2011}
}
"""






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



		def 						UpperCAmelCase							(					self)       ->       List[Any]:






							'''simple docstring'''
							return datasets.MetricInfo(
							    description=_DESCRIPTION				,		citation=_CITATION				,		inputs_description=_KWARGS_DESCRIPTION				,		features=datasets.Features(
							        {
							            '''predictions''': datasets.Value('''int32'''),
							            '''references''': datasets.Value('''int32'''),
							        })				,		reference_urls=[
							        '''https://scikit-learn.org/stable/modules/generated/sklearn.metrics.matthews_corrcoef.html'''
							    ]				,		)




		def 						UpperCAmelCase							(					self				,		__a				,		__a				,		__a=None)       ->       Dict:






							'''simple docstring'''
							return {
							    "matthews_correlation": float(matthews_corrcoef(__a				,		__a				,		sample_weight=__a)),
							}






