


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__					(     a__:						List[str]			)    ->						Any:


							'''simple docstring'''





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


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

def 			lowerCAmelCase__					(     a__:						str     ,    a__:						List[str]			)    ->						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}'
							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=a__     ,    filename=a__     ,    monitor=F'''val_{metric}'''     ,    mode='max'     ,    save_top_k=3     ,    every_n_epochs=1     ,    )
							return checkpoint_callback

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


							'''simple docstring'''





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







class 						__a (							pl.Callback     ):
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Tuple:




									"""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(_SCREAMING_SNAKE_CASE		)
		@rank_zero_only
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=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=_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
																							_UpperCAmelCase											=     metrics[key]
																							if isinstance(_SCREAMING_SNAKE_CASE     ,       torch.Tensor		):
																														_UpperCAmelCase											=     val.item()
																							_UpperCAmelCase											=     f'''{key}: {val:.6f}\n'''
																							writer.write(_SCREAMING_SNAKE_CASE		)

									if not save_generations:
																return

									if "preds" in metrics:
																_UpperCAmelCase											=     '\n'.join(metrics['preds']		)
																generations_file.open('w+'		).write(_SCREAMING_SNAKE_CASE		)
		@rank_zero_only
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Optional[Any]:




									"""simple docstring"""





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

									_UpperCAmelCase											=     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       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Optional[int]:




									"""simple docstring"""





									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       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			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")




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 AutoFeatureExtractor, WavaVecaFeatureExtractor
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_feature_extraction import CustomFeatureExtractor  # noqa E402


lowerCAmelCase__    :Optional[Any]       =				get_tests_dir('''fixtures''')

class 						__a (							unittest.TestCase     ):
		def       UpperCAmelCase__							(     self		)					->			Any:




									"""simple docstring"""





									_UpperCAmelCase											=     mock.Mock()
									_UpperCAmelCase											=     500
									_UpperCAmelCase											=     {}
									_UpperCAmelCase											=     HTTPError
									_UpperCAmelCase											=     {}

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






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




									"""simple docstring"""





									_UpperCAmelCase											=     WavaVecaFeatureExtractor.from_pretrained(
									    'https://huggingface.co/hf-internal-testing/tiny-random-wav2vec2/resolve/main/preprocessor_config.json'		)



@is_staging_test
class 						__a (							unittest.TestCase     ):
		@classmethod
		def       UpperCAmelCase__							(     cls		)					->			Tuple:




									"""simple docstring"""





									_UpperCAmelCase											=     TOKEN
									HfFolder.save_token(_SCREAMING_SNAKE_CASE		)
		@classmethod
		def       UpperCAmelCase__							(     cls		)					->			Tuple:




									"""simple docstring"""





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

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

									try:
																delete_repo(token=cls._token     ,       repo_id='test-dynamic-feature-extractor'		)
									except HTTPError:
																pass
		def       UpperCAmelCase__							(     self		)					->			Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase											=     WavaVecaFeatureExtractor.from_pretrained(_SCREAMING_SNAKE_CASE		)
									feature_extractor.push_to_hub('test-feature-extractor'     ,       use_auth_token=self._token		)

									_UpperCAmelCase											=     WavaVecaFeatureExtractor.from_pretrained(f'''{USER}/test-feature-extractor'''		)
									for k, v in feature_extractor.__dict__.items():
																self.assertEqual(_SCREAMING_SNAKE_CASE     ,       getattr(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)		)

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

									# Push to hub via save_pretrained
									with tempfile.TemporaryDirectory() as tmp_dir:
																feature_extractor.save_pretrained(
																    _SCREAMING_SNAKE_CASE     ,       repo_id='test-feature-extractor'     ,       push_to_hub=_SCREAMING_SNAKE_CASE     ,       use_auth_token=self._token		)

									_UpperCAmelCase											=     WavaVecaFeatureExtractor.from_pretrained(f'''{USER}/test-feature-extractor'''		)
									for k, v in feature_extractor.__dict__.items():
																self.assertEqual(_SCREAMING_SNAKE_CASE     ,       getattr(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)		)
		def       UpperCAmelCase__							(     self		)					->			Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     WavaVecaFeatureExtractor.from_pretrained(_SCREAMING_SNAKE_CASE		)
									feature_extractor.push_to_hub('valid_org/test-feature-extractor'     ,       use_auth_token=self._token		)

									_UpperCAmelCase											=     WavaVecaFeatureExtractor.from_pretrained('valid_org/test-feature-extractor'		)
									for k, v in feature_extractor.__dict__.items():
																self.assertEqual(_SCREAMING_SNAKE_CASE     ,       getattr(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)		)

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

									# Push to hub via save_pretrained
									with tempfile.TemporaryDirectory() as tmp_dir:
																feature_extractor.save_pretrained(
																    _SCREAMING_SNAKE_CASE     ,       repo_id='valid_org/test-feature-extractor-org'     ,       push_to_hub=_SCREAMING_SNAKE_CASE     ,       use_auth_token=self._token		)

									_UpperCAmelCase											=     WavaVecaFeatureExtractor.from_pretrained('valid_org/test-feature-extractor-org'		)
									for k, v in feature_extractor.__dict__.items():
																self.assertEqual(_SCREAMING_SNAKE_CASE     ,       getattr(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)		)






		def       UpperCAmelCase__							(     self		)					->			Tuple:




									"""simple docstring"""





									CustomFeatureExtractor.register_for_auto_class()
									_UpperCAmelCase											=     CustomFeatureExtractor.from_pretrained(_SCREAMING_SNAKE_CASE		)

									feature_extractor.push_to_hub('test-dynamic-feature-extractor'     ,       use_auth_token=self._token		)

									# This has added the proper auto_map field to the config
									self.assertDictEqual(
									    feature_extractor.auto_map     ,       {'AutoFeatureExtractor': 'custom_feature_extraction.CustomFeatureExtractor'}     ,       )

									_UpperCAmelCase											=     AutoFeatureExtractor.from_pretrained(
									    f'''{USER}/test-dynamic-feature-extractor'''     ,       trust_remote_code=_SCREAMING_SNAKE_CASE		)
									# Can't make an isinstance check because the new_feature_extractor is from the CustomFeatureExtractor class of a dynamic module
									self.assertEqual(new_feature_extractor.__class__.__name__     ,       'CustomFeatureExtractor'		)
