



"""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							(	__A      :						Dict     )     ->       Union[str, Any]:






				'''simple docstring'''
				SCREAMING_SNAKE_CASE						:							Optional[Any]							=     filter(lambda __A     : p.requires_grad							,		model.parameters()     )
				SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     sum([np.prod(p.size()     ) for p in model_parameters]     )
				return params


A_						:    Any				=      logging.getLogger(__name__)
def  __snake_case							(	__A      :						Union[str, Any]							,		__A      :						Optional[int]     )     ->       Any:






				'''simple docstring'''
				if metric == "rouge2":
								SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     '{val_avg_rouge2:.4f}-{step_count}'
				elif metric == "bleu":
								SCREAMING_SNAKE_CASE						:							Optional[Any]							=     '{val_avg_bleu:.4f}-{step_count}'
				elif metric == "em":
								SCREAMING_SNAKE_CASE						:							List[Any]							=     '{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.'     )

				SCREAMING_SNAKE_CASE						:							Optional[int]							=     ModelCheckpoint(
				    dirpath=__A							,		filename=__A							,		monitor=F"""val_{metric}"""							,		mode='max'							,		save_top_k=3							,		every_n_epochs=1							,		)
				return checkpoint_callback
def  __snake_case							(	__A      :						int							,		__A      :						Optional[int]     )     ->       str:






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







class     lowerCAmelCase__						(	pl.Callback ):


	'''simple docstring'''







	def 			_lowerCAmelCase     (      self				:       Tuple     ,       _SCREAMING_SNAKE_CASE				:       Union[str, Any]     ,       _SCREAMING_SNAKE_CASE				:       Union[str, Any]    )						->	Any:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Dict							=     {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 			_lowerCAmelCase     (      self				:       Optional[int]     ,       _SCREAMING_SNAKE_CASE				:       pl.Trainer     ,       _SCREAMING_SNAKE_CASE				:       pl.LightningModule     ,       _SCREAMING_SNAKE_CASE				:       str     ,       _SCREAMING_SNAKE_CASE				:       Optional[int]=True    )						->	None:





					"""simple docstring"""


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

					if not save_generations:
									return

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





	@rank_zero_only
	def 			_lowerCAmelCase     (      self				:       Any     ,       _SCREAMING_SNAKE_CASE				:       Any     ,       _SCREAMING_SNAKE_CASE				:       Any    )						->	List[Any]:





					"""simple docstring"""


					try:
									SCREAMING_SNAKE_CASE						:							Optional[int]							=     pl_module.model.model.num_parameters()
					except AttributeError:
									SCREAMING_SNAKE_CASE						:							Tuple							=     pl_module.model.num_parameters()

					SCREAMING_SNAKE_CASE						:							List[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 			_lowerCAmelCase     (      self				:       Optional[Any]     ,       _SCREAMING_SNAKE_CASE				:       pl.Trainer     ,       _SCREAMING_SNAKE_CASE				:       pl.LightningModule    )						->	List[str]:





					"""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 			_lowerCAmelCase     (      self				:       int     ,       _SCREAMING_SNAKE_CASE				:       pl.Trainer     ,       _SCREAMING_SNAKE_CASE				:       Optional[Any]    )						->	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"""

import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    FEATURE_EXTRACTOR_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


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

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


A_						:    int				=      get_tests_dir('fixtures')
A_						:    int				=      get_tests_dir('fixtures/dummy_feature_extractor_config.json')
A_						:    Dict				=      get_tests_dir('fixtures/dummy-config.json')
class     lowerCAmelCase__						(	unittest.TestCase ):


	'''simple docstring'''







	def 			_lowerCAmelCase     (      self				:       int    )						->	Tuple:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Tuple							=     0





	def 			_lowerCAmelCase     (      self				:       List[str]    )						->	Tuple:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Any							=     AutoFeatureExtractor.from_pretrained('facebook/wav2vec2-base-960h'    )
					self.assertIsInstance(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE    )





	def 			_lowerCAmelCase     (      self				:       Dict    )						->	Dict:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							str							=     AutoFeatureExtractor.from_pretrained(_SCREAMING_SNAKE_CASE    )
					self.assertIsInstance(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE    )





	def 			_lowerCAmelCase     (      self				:       Optional[Any]    )						->	Optional[int]:





					"""simple docstring"""


					with tempfile.TemporaryDirectory() as tmpdirname:
									SCREAMING_SNAKE_CASE						:							List[Any]							=     WavaVecaConfig()

									# remove feature_extractor_type to make sure config.json alone is enough to load feature processor locally
									SCREAMING_SNAKE_CASE						:							Optional[Any]							=     AutoFeatureExtractor.from_pretrained(_SCREAMING_SNAKE_CASE    ).to_dict()

									config_dict.pop('feature_extractor_type'    )
									SCREAMING_SNAKE_CASE						:							int							=     WavaVecaFeatureExtractor(**_SCREAMING_SNAKE_CASE    )

									# save in new folder
									model_config.save_pretrained(_SCREAMING_SNAKE_CASE    )
									config.save_pretrained(_SCREAMING_SNAKE_CASE    )

									SCREAMING_SNAKE_CASE						:							List[Any]							=     AutoFeatureExtractor.from_pretrained(_SCREAMING_SNAKE_CASE    )

									# make sure private variable is not incorrectly saved
									SCREAMING_SNAKE_CASE						:							Optional[int]							=     json.loads(config.to_json_string()    )
									self.assertTrue('_processor_class' not in dict_as_saved    )

					self.assertIsInstance(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE    )





	def 			_lowerCAmelCase     (      self				:       List[str]    )						->	int:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							List[Any]							=     AutoFeatureExtractor.from_pretrained(_SCREAMING_SNAKE_CASE    )
					self.assertIsInstance(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE    )





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





					"""simple docstring"""


					with self.assertRaisesRegex(
					    _SCREAMING_SNAKE_CASE     ,       'bert-base is not a local folder and is not a valid model identifier'    ):
									SCREAMING_SNAKE_CASE						:							str							=     AutoFeatureExtractor.from_pretrained('bert-base'    )





	def 			_lowerCAmelCase     (      self				:       str    )						->	str:





					"""simple docstring"""


					with self.assertRaisesRegex(
					    _SCREAMING_SNAKE_CASE     ,       R'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'    ):
									SCREAMING_SNAKE_CASE						:							List[str]							=     AutoFeatureExtractor.from_pretrained(_SCREAMING_SNAKE_CASE     ,       revision='aaaaaa'    )





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





					"""simple docstring"""


					with self.assertRaisesRegex(
					    _SCREAMING_SNAKE_CASE     ,       'hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json.'     ,       ):
									SCREAMING_SNAKE_CASE						:							int							=     AutoFeatureExtractor.from_pretrained('hf-internal-testing/config-no-model'    )





	def 			_lowerCAmelCase     (      self				:       Optional[Any]    )						->	str:





					"""simple docstring"""


					with self.assertRaises(_SCREAMING_SNAKE_CASE    ):
									SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     AutoFeatureExtractor.from_pretrained(
									    'hf-internal-testing/test_dynamic_feature_extractor'    )
					# If remote code is disabled, we can't load this config.
					with self.assertRaises(_SCREAMING_SNAKE_CASE    ):
									SCREAMING_SNAKE_CASE						:							int							=     AutoFeatureExtractor.from_pretrained(
									    'hf-internal-testing/test_dynamic_feature_extractor'     ,       trust_remote_code=_SCREAMING_SNAKE_CASE    )

					SCREAMING_SNAKE_CASE						:							Optional[int]							=     AutoFeatureExtractor.from_pretrained(
					    'hf-internal-testing/test_dynamic_feature_extractor'     ,       trust_remote_code=_SCREAMING_SNAKE_CASE    )
					self.assertEqual(feature_extractor.__class__.__name__     ,       'NewFeatureExtractor'    )

					# Test feature extractor can be reloaded.
					with tempfile.TemporaryDirectory() as tmp_dir:
									feature_extractor.save_pretrained(_SCREAMING_SNAKE_CASE    )
									SCREAMING_SNAKE_CASE						:							List[str]							=     AutoFeatureExtractor.from_pretrained(_SCREAMING_SNAKE_CASE     ,       trust_remote_code=_SCREAMING_SNAKE_CASE    )
					self.assertEqual(reloaded_feature_extractor.__class__.__name__     ,       'NewFeatureExtractor'    )





	def 			_lowerCAmelCase     (      self				:       Any    )						->	Optional[int]:





					"""simple docstring"""


					try:
									AutoConfig.register('custom'     ,       _SCREAMING_SNAKE_CASE    )
									AutoFeatureExtractor.register(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE    )
									# Trying to register something existing in the Transformers library will raise an error
									with self.assertRaises(_SCREAMING_SNAKE_CASE    ):
													AutoFeatureExtractor.register(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE    )

									# Now that the config is registered, it can be used as any other config with the auto-API
									SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     CustomFeatureExtractor.from_pretrained(_SCREAMING_SNAKE_CASE    )
									with tempfile.TemporaryDirectory() as tmp_dir:
													feature_extractor.save_pretrained(_SCREAMING_SNAKE_CASE    )
													SCREAMING_SNAKE_CASE						:							Dict							=     AutoFeatureExtractor.from_pretrained(_SCREAMING_SNAKE_CASE    )
													self.assertIsInstance(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE    )

					finally:
									if "custom" in CONFIG_MAPPING._extra_content:
													del CONFIG_MAPPING._extra_content["custom"]
									if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
													del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]







	def 			_lowerCAmelCase     (      self				:       Tuple    )						->	str:





					"""simple docstring"""


					class     lowerCAmelCase__						(	_lowerCamelCase ):


						'''simple docstring'''







						_SCREAMING_SNAKE_CASE	:		str  						= True

					try:
									AutoConfig.register('custom'     ,       _SCREAMING_SNAKE_CASE    )
									AutoFeatureExtractor.register(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE    )
									# If remote code is not set, the default is to use local
									SCREAMING_SNAKE_CASE						:							str							=     AutoFeatureExtractor.from_pretrained(
									    'hf-internal-testing/test_dynamic_feature_extractor'    )
									self.assertEqual(feature_extractor.__class__.__name__     ,       'NewFeatureExtractor'    )
									self.assertTrue(feature_extractor.is_local    )

									# If remote code is disabled, we load the local one.
									SCREAMING_SNAKE_CASE						:							List[Any]							=     AutoFeatureExtractor.from_pretrained(
									    'hf-internal-testing/test_dynamic_feature_extractor'     ,       trust_remote_code=_SCREAMING_SNAKE_CASE    )
									self.assertEqual(feature_extractor.__class__.__name__     ,       'NewFeatureExtractor'    )
									self.assertTrue(feature_extractor.is_local    )

									# If remote is enabled, we load from the Hub
									SCREAMING_SNAKE_CASE						:							Optional[int]							=     AutoFeatureExtractor.from_pretrained(
									    'hf-internal-testing/test_dynamic_feature_extractor'     ,       trust_remote_code=_SCREAMING_SNAKE_CASE    )
									self.assertEqual(feature_extractor.__class__.__name__     ,       'NewFeatureExtractor'    )
									self.assertTrue(not hasattr(_SCREAMING_SNAKE_CASE     ,       'is_local'    )    )

					finally:
									if "custom" in CONFIG_MAPPING._extra_content:
													del CONFIG_MAPPING._extra_content["custom"]
									if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
													del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]





