



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


snake_case   					=				logging.getLogger(__name__)




def 			snake_case		(		lowerCAmelCase_  ,						lowerCAmelCase_	)						->				Optional[int]:
  if metric == "rouge2":
    _snake_case   	=      '''{val_avg_rouge2:.4f}-{step_count}'''
  elif metric == "bleu":
    _snake_case   	=      '''{val_avg_bleu:.4f}-{step_count}'''
  elif metric == "em":
    _snake_case   	=      '''{val_avg_em:.4f}-{step_count}'''
  elif metric == "loss":
    _snake_case   	=      '''{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.'''	)

  _snake_case   	=      ModelCheckpoint(
      dirpath=lowerCAmelCase_  ,						filename=lowerCAmelCase_  ,						monitor=f"""val_{metric}"""  ,						mode='''max'''  ,						save_top_k=1  ,						every_n_epochs=1  ,						)

  return checkpoint_callback




def 			snake_case		(		lowerCAmelCase_  ,						lowerCAmelCase_	)						->				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    __UpperCAmelCase    (       self    :	Optional[int]							,       __lowerCamelCase    :	Union[str, Any]							,       __lowerCamelCase    :	Dict				):



     """simple docstring"""
     _snake_case   	=      {f"""lr_group_{i}""": param['''lr'''] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups				)}
     pl_module.logger.log_metrics(__lowerCamelCase				)



   @rank_zero_only
   def    __UpperCAmelCase    (       self    :	Optional[int]							,       __lowerCamelCase    :	pl.Trainer							,       __lowerCamelCase    :	pl.LightningModule							,       __lowerCamelCase    :	str							,       __lowerCamelCase    :	str=True				):



     """simple docstring"""
     logger.info(f"""***** {type_path} results at step {trainer.global_step:05d} *****"""				)
     _snake_case   	=      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   	=      Path(pl_module.hparams.output_dir				)
     if type_path == "test":
       _snake_case   	=      od / '''test_results.txt'''
       _snake_case   	=      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   	=      od / f"""{type_path}_results/{trainer.global_step:05d}.txt"""
       _snake_case   	=      od / f"""{type_path}_generations/{trainer.global_step:05d}.txt"""
       results_file.parent.mkdir(exist_ok=__lowerCamelCase				)
       generations_file.parent.mkdir(exist_ok=__lowerCamelCase				)
     with open(__lowerCamelCase							,       '''a+'''				) as writer:
       for key in sorted(__lowerCamelCase				):
         if key in ["log", "progress_bar", "preds"]:
           continue
         _snake_case   	=      metrics[key]
         if isinstance(__lowerCamelCase							,       torch.Tensor				):
           _snake_case   	=      val.item()
         _snake_case   	=      f"""{key}: {val:.6f}\n"""
         writer.write(__lowerCamelCase				)

     if not save_generations:
       return

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



   @rank_zero_only
   def    __UpperCAmelCase    (       self    :	str							,       __lowerCamelCase    :	Optional[int]							,       __lowerCamelCase    :	str				):



     """simple docstring"""
     try:
       _snake_case   	=      pl_module.model.model.num_parameters()
     except AttributeError:
       _snake_case   	=      pl_module.model.num_parameters()

     _snake_case   	=      count_trainable_parameters(__lowerCamelCase				)
     # 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    :	str							,       __lowerCamelCase    :	pl.Trainer							,       __lowerCamelCase    :	pl.LightningModule				):



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





   @rank_zero_only
   def    __UpperCAmelCase    (       self    :	Any							,       __lowerCamelCase    :	pl.Trainer							,       __lowerCamelCase    :	List[Any]				):



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

import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoImageProcessor, ViTImageProcessor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / 'utils'))

from test_module.custom_image_processing import CustomImageProcessor  # noqa E402


UpperCamelCase__  :						str					=      get_tests_dir('fixtures')


class 					_lowercase						(   unittest.TestCase ):
      '''simple docstring'''


      def 		lowerCAmelCase__       (					self				)      ->      int:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[int] 					=	mock.Mock()
         UpperCAmelCase__				:					Dict 					=	500
         UpperCAmelCase__				:					Optional[int] 					=	{}
         UpperCAmelCase__				:					List[str] 					=	HTTPError
         UpperCAmelCase__				:					Tuple 					=	{}

         # Download this model to make sure it's in the cache.
         UpperCAmelCase__				:					Any 					=	ViTImageProcessor.from_pretrained('''hf-internal-testing/tiny-random-vit'''				)
         # Under the mock environment we get a 500 error when trying to reach the model.
         with mock.patch('''requests.Session.request'''       ,return_value=lowerCamelCase_				) as mock_head:
            UpperCAmelCase__				:					Tuple 					=	ViTImageProcessor.from_pretrained('''hf-internal-testing/tiny-random-vit'''				)
            # This check we did call the fake head request
            mock_head.assert_called()




      def 		lowerCAmelCase__       (					self				)      ->      Tuple:



         '''simple docstring'''



         UpperCAmelCase__				:					Dict 					=	ViTImageProcessor.from_pretrained(
             '''https://huggingface.co/hf-internal-testing/tiny-random-vit/resolve/main/preprocessor_config.json'''				)




      def 		lowerCAmelCase__       (					self				)      ->      Union[str, Any]:



         '''simple docstring'''



         with self.assertRaises(lowerCamelCase_				):
            # config is in subfolder, the following should not work without specifying the subfolder
            UpperCAmelCase__				:					int 					=	AutoImageProcessor.from_pretrained('''hf-internal-testing/stable-diffusion-all-variants'''				)

         UpperCAmelCase__				:					Optional[Any] 					=	AutoImageProcessor.from_pretrained(
             '''hf-internal-testing/stable-diffusion-all-variants'''       ,subfolder='''feature_extractor'''				)

         self.assertIsNotNone(lowerCamelCase_				)




@is_staging_test
class 					_lowercase						(   unittest.TestCase ):
      '''simple docstring'''


      @classmethod
      def 		lowerCAmelCase__       (					cls				)      ->      Optional[Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[int] 					=	TOKEN
         HfFolder.save_token(lowerCamelCase_				)




      @classmethod
      def 		lowerCAmelCase__       (					cls				)      ->      List[str]:



         '''simple docstring'''



         try:
            delete_repo(token=cls._token       ,repo_id='''test-image-processor'''				)
         except HTTPError:
            pass

         try:
            delete_repo(token=cls._token       ,repo_id='''valid_org/test-image-processor-org'''				)
         except HTTPError:
            pass

         try:
            delete_repo(token=cls._token       ,repo_id='''test-dynamic-image-processor'''				)
         except HTTPError:
            pass




      def 		lowerCAmelCase__       (					self				)      ->      Union[str, Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[Any] 					=	ViTImageProcessor.from_pretrained(lowerCamelCase_				)
         image_processor.push_to_hub('''test-image-processor'''       ,use_auth_token=self._token				)

         UpperCAmelCase__				:					Any 					=	ViTImageProcessor.from_pretrained(f'''{USER}/test-image-processor'''				)
         for k, v in image_processor.__dict__.items():
            self.assertEqual(lowerCamelCase_       ,getattr(lowerCamelCase_       ,lowerCamelCase_				)				)

         # Reset repo
         delete_repo(token=self._token       ,repo_id='''test-image-processor'''				)

         # Push to hub via save_pretrained
         with tempfile.TemporaryDirectory() as tmp_dir:
            image_processor.save_pretrained(
                lowerCamelCase_       ,repo_id='''test-image-processor'''       ,push_to_hub=lowerCamelCase_       ,use_auth_token=self._token				)

         UpperCAmelCase__				:					List[Any] 					=	ViTImageProcessor.from_pretrained(f'''{USER}/test-image-processor'''				)
         for k, v in image_processor.__dict__.items():
            self.assertEqual(lowerCamelCase_       ,getattr(lowerCamelCase_       ,lowerCamelCase_				)				)




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



         '''simple docstring'''



         UpperCAmelCase__				:					Union[str, Any] 					=	ViTImageProcessor.from_pretrained(lowerCamelCase_				)
         image_processor.push_to_hub('''valid_org/test-image-processor'''       ,use_auth_token=self._token				)

         UpperCAmelCase__				:					Any 					=	ViTImageProcessor.from_pretrained('''valid_org/test-image-processor'''				)
         for k, v in image_processor.__dict__.items():
            self.assertEqual(lowerCamelCase_       ,getattr(lowerCamelCase_       ,lowerCamelCase_				)				)

         # Reset repo
         delete_repo(token=self._token       ,repo_id='''valid_org/test-image-processor'''				)

         # Push to hub via save_pretrained
         with tempfile.TemporaryDirectory() as tmp_dir:
            image_processor.save_pretrained(
                lowerCamelCase_       ,repo_id='''valid_org/test-image-processor-org'''       ,push_to_hub=lowerCamelCase_       ,use_auth_token=self._token				)

         UpperCAmelCase__				:					Any 					=	ViTImageProcessor.from_pretrained('''valid_org/test-image-processor-org'''				)
         for k, v in image_processor.__dict__.items():
            self.assertEqual(lowerCamelCase_       ,getattr(lowerCamelCase_       ,lowerCamelCase_				)				)




      def 		lowerCAmelCase__       (					self				)      ->      str:



         '''simple docstring'''



         CustomImageProcessor.register_for_auto_class()
         UpperCAmelCase__				:					List[str] 					=	CustomImageProcessor.from_pretrained(lowerCamelCase_				)

         image_processor.push_to_hub('''test-dynamic-image-processor'''       ,use_auth_token=self._token				)

         # This has added the proper auto_map field to the config
         self.assertDictEqual(
             image_processor.auto_map       ,{'''AutoImageProcessor''': '''custom_image_processing.CustomImageProcessor'''}       ,)

         UpperCAmelCase__				:					Any 					=	AutoImageProcessor.from_pretrained(
             f'''{USER}/test-dynamic-image-processor'''       ,trust_remote_code=lowerCamelCase_				)
         # Can't make an isinstance check because the new_image_processor is from the CustomImageProcessor class of a dynamic module
         self.assertEqual(new_image_processor.__class__.__name__       ,'''CustomImageProcessor'''				)




