


# tests directory-specific settings - this file is run automatically
# by pytest before any tests are run

import doctest
import sys
import warnings
from os.path import abspath, dirname, join

import _pytest

from transformers.testing_utils import HfDoctestModule, HfDocTestParser


# allow having multiple repository checkouts and not needing to remember to rerun
# 'pip install -e .[dev]' when switching between checkouts and running tests.
UpperCAmelCase_				:			Optional[int]   =	abspath(join(dirname(__file__), 'src'))
sys.path.insert(1, git_repo_path)

# silence FutureWarning warnings in tests since often we can't act on them until
# they become normal warnings - i.e. the tests still need to test the current functionality
warnings.simplefilter(action='ignore', category=FutureWarning)







def 	SCREAMING_SNAKE_CASE_						(      __A			:      List[Any]    )     ->		Optional[Any]:

       """simple docstring"""

       config.addinivalue_line(
           'markers'  ,   'is_pt_tf_cross_test: mark test to run only when PT and TF interactions are tested'    )
       config.addinivalue_line(
           'markers'  ,   'is_pt_flax_cross_test: mark test to run only when PT and FLAX interactions are tested'    )
       config.addinivalue_line('markers'  ,   'is_pipeline_test: mark test to run only when pipelines are tested'    )
       config.addinivalue_line('markers'  ,   'is_staging_test: mark test to run only in the staging environment'    )
       config.addinivalue_line('markers'  ,   'accelerate_tests: mark test that require accelerate'    )
       config.addinivalue_line('markers'  ,   'tool_tests: mark the tool tests that are run on their specific schedule'    )







def 	SCREAMING_SNAKE_CASE_						(      __A			:      Tuple    )     ->		List[Any]:

       """simple docstring"""

       from transformers.testing_utils import pytest_addoption_shared

       pytest_addoption_shared(__A    )







def 	SCREAMING_SNAKE_CASE_						(      __A			:      int    )     ->		List[str]:

       """simple docstring"""

       from transformers.testing_utils import pytest_terminal_summary_main

       a_ :		Optional[int]							=  terminalreporter.config.getoption('--make-reports'    )
       if make_reports:
              pytest_terminal_summary_main(__A  ,   id=__A    )







def 	SCREAMING_SNAKE_CASE_						(      __A			:      List[str]  ,   __A			:      Optional[int]    )     ->		List[str]:

       """simple docstring"""

       if exitstatus == 5:
              a_ :		Optional[Any]							=  0


# Doctest custom flag to ignore output.
UpperCAmelCase_				:			int   =	doctest.register_optionflag('IGNORE_RESULT')

UpperCAmelCase_				:			List[str]   =	doctest.OutputChecker
class 			SCREAMING_SNAKE_CASE__  (    lowercase__      ):
      def 	SCREAMING_SNAKE_CASE							(	self    :	str					,							SCREAMING_SNAKE_CASE__    :	Optional[Any]					,							SCREAMING_SNAKE_CASE__    :	Tuple					,							SCREAMING_SNAKE_CASE__    :	Union[str, Any]					)		->      List[Any]:
             if IGNORE_RESULT & optionflags:
                    return True
             return OutputChecker.check_output(self					,							SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					)
UpperCAmelCase_				:			Optional[int]   =	CustomOutputChecker
UpperCAmelCase_				:			Optional[Any]   =	HfDoctestModule
UpperCAmelCase_				:			Optional[int]   =	HfDocTestParser






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 	SCREAMING_SNAKE_CASE_						(      __A			:      Any    )     ->		Optional[int]:

       """simple docstring"""

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


UpperCAmelCase_				:			Dict   =	logging.getLogger(__name__)







def 	SCREAMING_SNAKE_CASE_						(      __A			:      int  ,   __A			:      int    )     ->		List[str]:

       """simple docstring"""

       if metric == "rouge2":
              a_ :		Dict							=  '{val_avg_rouge2:.4f}-{step_count}'
       elif metric == "bleu":
              a_ :		Dict							=  '{val_avg_bleu:.4f}-{step_count}'
       elif metric == "em":
              a_ :		Union[str, Any]							=  '{val_avg_em:.4f}-{step_count}'
       elif metric == "loss":
              a_ :		str							=  '{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.'    )

       a_ :		Dict							=  ModelCheckpoint(
           dirpath=__A  ,   filename=__A  ,   monitor=F"""val_{metric}"""  ,   mode='max'  ,   save_top_k=1  ,   every_n_epochs=1  ,   )

       return checkpoint_callback







def 	SCREAMING_SNAKE_CASE_						(      __A			:      int  ,   __A			:      List[Any]    )     ->		int:

       """simple docstring"""

       return EarlyStopping(
           monitor=F"""val_{metric}"""  ,   mode='min' if 'loss' in metric else 'max'  ,   patience=__A  ,   verbose=__A  ,   )
class 			SCREAMING_SNAKE_CASE__  (    pl.Callback      ):
      def 	SCREAMING_SNAKE_CASE							(	self    :	Dict					,							SCREAMING_SNAKE_CASE__    :	Union[str, Any]					,							SCREAMING_SNAKE_CASE__    :	List[str]					)		->      Any:
             a_ :		int							=  {F"""lr_group_{i}""": param['lr'] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups					)}
             pl_module.logger.log_metrics(SCREAMING_SNAKE_CASE__					)
      @rank_zero_only
      def 	SCREAMING_SNAKE_CASE							(	self    :	List[str]					,							SCREAMING_SNAKE_CASE__    :	pl.Trainer					,							SCREAMING_SNAKE_CASE__    :	pl.LightningModule					,							SCREAMING_SNAKE_CASE__    :	str					,							SCREAMING_SNAKE_CASE__    :	List[str]=True					)		->      None:
             logger.info(F"""***** {type_path} results at step {trainer.global_step:05d} *****"""					)
             a_ :		List[str]							=  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
             a_ :		Optional[int]							=  Path(pl_module.hparams.output_dir					)
             if type_path == "test":
                    a_ :		Tuple							=  od / 'test_results.txt'
                    a_ :		Optional[int]							=  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.
                    a_ :		str							=  od / F"""{type_path}_results/{trainer.global_step:05d}.txt"""
                    a_ :		Optional[Any]							=  od / F"""{type_path}_generations/{trainer.global_step:05d}.txt"""
                    results_file.parent.mkdir(exist_ok=SCREAMING_SNAKE_CASE__					)
                    generations_file.parent.mkdir(exist_ok=SCREAMING_SNAKE_CASE__					)
             with open(SCREAMING_SNAKE_CASE__					,							'a+'					) as writer:
                    for key in sorted(SCREAMING_SNAKE_CASE__					):
                           if key in ["log", "progress_bar", "preds"]:
                                  continue
                           a_ :		int							=  metrics[key]
                           if isinstance(SCREAMING_SNAKE_CASE__					,							torch.Tensor					):
                                  a_ :		List[str]							=  val.item()
                           a_ :		int							=  F"""{key}: {val:.6f}\n"""
                           writer.write(SCREAMING_SNAKE_CASE__					)

             if not save_generations:
                    return

             if "preds" in metrics:
                    a_ :		Optional[Any]							=  '\n'.join(metrics['preds']					)
                    generations_file.open('w+'					).write(SCREAMING_SNAKE_CASE__					)
      @rank_zero_only
      def 	SCREAMING_SNAKE_CASE							(	self    :	Any					,							SCREAMING_SNAKE_CASE__    :	Dict					,							SCREAMING_SNAKE_CASE__    :	Union[str, Any]					)		->      Dict:
             try:
                    a_ :		Any							=  pl_module.model.model.num_parameters()
             except AttributeError:
                    a_ :		List[str]							=  pl_module.model.num_parameters()

             a_ :		Any							=  count_trainable_parameters(SCREAMING_SNAKE_CASE__					)
             # 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 	SCREAMING_SNAKE_CASE							(	self    :	Union[str, Any]					,							SCREAMING_SNAKE_CASE__    :	pl.Trainer					,							SCREAMING_SNAKE_CASE__    :	pl.LightningModule					)		->      Optional[int]:
             save_json(pl_module.metrics					,							pl_module.metrics_save_path					)
             return self._write_logs(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							'test'					)

      @rank_zero_only
      def 	SCREAMING_SNAKE_CASE							(	self    :	Any					,							SCREAMING_SNAKE_CASE__    :	pl.Trainer					,							SCREAMING_SNAKE_CASE__    :	Optional[Any]					)		->      Optional[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")


