



import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
							from transformers import WhisperFeatureExtractor

if is_torch_available():
							import torch

UpperCAmelCase__     = random.Random()





def 							A			(  _UpperCAmelCase       :     Optional[Any]							,     _UpperCAmelCase       :     Dict=1.0							,     _UpperCAmelCase       :     List[str]=None							,     _UpperCAmelCase       :     List[str]=None					)       ->				List[Any]:







			'''simple docstring'''
			if rng is None:
						_UpperCAmelCase	 =     global_rng

			_UpperCAmelCase	 =     []
			for batch_idx in range(shape[0]					):
						values.append([]					)
						for _ in range(shape[1]					):
									values[-1].append(rng.random() * scale					)

			return values







@require_torch
@require_torchaudio
class     __lowerCAmelCase					(    unittest.TestCase		):



			def __init__(		self     :						Tuple				,	A     :						Dict				,	A     :						Optional[Any]=7				,	A     :						Optional[Any]=4_00				,	A     :						List[str]=20_00				,	A     :						Tuple=10				,	A     :						Tuple=1_60				,	A     :						Tuple=8				,	A     :						int=0.0				,	A     :						Optional[Any]=40_00				,	A     :						List[Any]=False				,	A     :						str=True				,	)     ->					str:



						"""simple docstring"""






						_UpperCAmelCase	 =     parent
						_UpperCAmelCase	 =     batch_size
						_UpperCAmelCase	 =     min_seq_length
						_UpperCAmelCase	 =     max_seq_length
						_UpperCAmelCase	 =     (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
						_UpperCAmelCase	 =     padding_value
						_UpperCAmelCase	 =     sampling_rate
						_UpperCAmelCase	 =     return_attention_mask
						_UpperCAmelCase	 =     do_normalize
						_UpperCAmelCase	 =     feature_size
						_UpperCAmelCase	 =     chunk_length
						_UpperCAmelCase	 =     hop_length



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



						"""simple docstring"""






						return {
						    "feature_size": self.feature_size,
						    "hop_length": self.hop_length,
						    "chunk_length": self.chunk_length,
						    "padding_value": self.padding_value,
						    "sampling_rate": self.sampling_rate,
						    "return_attention_mask": self.return_attention_mask,
						    "do_normalize": self.do_normalize,
						}



			def 					_lowerCamelCase (		self     :						Tuple				,	A     :						Tuple=False				,	A     :						Optional[Any]=False)     ->					Optional[Any]:



						"""simple docstring"""






						def _flatten(A     :						Union[str, Any]):
									return list(itertools.chain(*A))

						if equal_length:
									_UpperCAmelCase	 =     [floats_list((self.max_seq_length, self.feature_size)) for _ in range(self.batch_size)]
						else:
									# make sure that inputs increase in size
									_UpperCAmelCase	 =     [
									    floats_list((x, self.feature_size))
									    for x in range(self.min_seq_length				,	self.max_seq_length				,	self.seq_length_diff)
									]
						if numpify:
									_UpperCAmelCase	 =     [np.asarray(A) for x in speech_inputs]
						return speech_inputs







@require_torch
@require_torchaudio
class     __lowerCAmelCase					(    A       ,		unittest.TestCase		):
			UpperCamelCase           =  WhisperFeatureExtractor if is_speech_available() else None



			def 					_lowerCamelCase (		self     :						List[Any])     ->					int:



						"""simple docstring"""






						_UpperCAmelCase	 =     WhisperFeatureExtractionTester(self)



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



						"""simple docstring"""






						_UpperCAmelCase	 =     self.feature_extraction_class(**self.feat_extract_dict)

						with tempfile.TemporaryDirectory() as tmpdirname:
									_UpperCAmelCase	 =     feat_extract_first.save_pretrained(A)[0]
									check_json_file_has_correct_format(A)
									_UpperCAmelCase	 =     self.feature_extraction_class.from_pretrained(A)

						_UpperCAmelCase	 =     feat_extract_first.to_dict()
						_UpperCAmelCase	 =     feat_extract_second.to_dict()
						_UpperCAmelCase	 =     feat_extract_first.mel_filters
						_UpperCAmelCase	 =     feat_extract_second.mel_filters
						self.assertTrue(np.allclose(A				,	A))
						self.assertEqual(A				,	A)



			def 					_lowerCamelCase (		self     :						str)     ->					str:



						"""simple docstring"""






						_UpperCAmelCase	 =     self.feature_extraction_class(**self.feat_extract_dict)

						with tempfile.TemporaryDirectory() as tmpdirname:
									_UpperCAmelCase	 =     os.path.join(A				,	'feat_extract.json')
									feat_extract_first.to_json_file(A)
									_UpperCAmelCase	 =     self.feature_extraction_class.from_json_file(A)

						_UpperCAmelCase	 =     feat_extract_first.to_dict()
						_UpperCAmelCase	 =     feat_extract_second.to_dict()
						_UpperCAmelCase	 =     feat_extract_first.mel_filters
						_UpperCAmelCase	 =     feat_extract_second.mel_filters
						self.assertTrue(np.allclose(A				,	A))
						self.assertEqual(A				,	A)



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



						"""simple docstring"""






						_UpperCAmelCase	 =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict())
						# create three inputs of length 800, 1000, and 1200
						_UpperCAmelCase	 =     [floats_list((1, x))[0] for x in range(8_00				,	14_00				,	2_00)]
						_UpperCAmelCase	 =     [np.asarray(A) for speech_input in speech_inputs]

						# Test feature size
						_UpperCAmelCase	 =     feature_extractor(A				,	padding='max_length'				,	return_tensors='np').input_features
						self.assertTrue(input_features.ndim == 3)
						self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames)
						self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size)

						# Test not batched input
						_UpperCAmelCase	 =     feature_extractor(speech_inputs[0]				,	return_tensors='np').input_features
						_UpperCAmelCase	 =     feature_extractor(np_speech_inputs[0]				,	return_tensors='np').input_features
						self.assertTrue(np.allclose(A				,	A				,	atol=1E-3))

						# Test batched
						_UpperCAmelCase	 =     feature_extractor(A				,	return_tensors='np').input_features
						_UpperCAmelCase	 =     feature_extractor(A				,	return_tensors='np').input_features
						for enc_seq_a, enc_seq_a in zip(A				,	A):
									self.assertTrue(np.allclose(A				,	A				,	atol=1E-3))

						# Test 2-D numpy arrays are batched.
						_UpperCAmelCase	 =     [floats_list((1, x))[0] for x in (8_00, 8_00, 8_00)]
						_UpperCAmelCase	 =     np.asarray(A)
						_UpperCAmelCase	 =     feature_extractor(A				,	return_tensors='np').input_features
						_UpperCAmelCase	 =     feature_extractor(A				,	return_tensors='np').input_features
						for enc_seq_a, enc_seq_a in zip(A				,	A):
									self.assertTrue(np.allclose(A				,	A				,	atol=1E-3))

						# Test truncation required
						_UpperCAmelCase	 =     [floats_list((1, x))[0] for x in range(2_00				,	(feature_extractor.n_samples + 5_00)				,	2_00)]
						_UpperCAmelCase	 =     [np.asarray(A) for speech_input in speech_inputs]

						_UpperCAmelCase	 =     [x[: feature_extractor.n_samples] for x in speech_inputs]
						_UpperCAmelCase	 =     [np.asarray(A) for speech_input in speech_inputs_truncated]

						_UpperCAmelCase	 =     feature_extractor(A				,	return_tensors='np').input_features
						_UpperCAmelCase	 =     feature_extractor(A				,	return_tensors='np').input_features
						for enc_seq_a, enc_seq_a in zip(A				,	A):
									self.assertTrue(np.allclose(A				,	A				,	atol=1E-3))



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



						"""simple docstring"""






						import torch

						_UpperCAmelCase	 =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict())
						_UpperCAmelCase	 =     np.random.rand(1_00				,	32).astype(np.floataa)
						_UpperCAmelCase	 =     np_speech_inputs.tolist()

						for inputs in [py_speech_inputs, np_speech_inputs]:
									_UpperCAmelCase	 =     feature_extractor.pad([{'input_features': inputs}]				,	return_tensors='np')
									self.assertTrue(np_processed.input_features.dtype == np.floataa)
									_UpperCAmelCase	 =     feature_extractor.pad([{'input_features': inputs}]				,	return_tensors='pt')
									self.assertTrue(pt_processed.input_features.dtype == torch.floataa)



			def 					_lowerCamelCase (		self     :						Any				,	A     :						str)     ->					List[str]:



						"""simple docstring"""






						_UpperCAmelCase	 =     load_dataset('hf-internal-testing/librispeech_asr_dummy'				,	'clean'				,	split='validation')
						# automatic decoding with librispeech
						_UpperCAmelCase	 =     ds.sort('id').select(range(A))[:num_samples]['audio']

						return [x["array"] for x in speech_samples]



			def 					_lowerCamelCase (		self     :						Optional[int])     ->					Tuple:



						"""simple docstring"""






						_UpperCAmelCase	 =     torch.tensor(
						    [
						        0.1_1_9_3, -0.0_9_4_6, -0.1_0_9_8, -0.0_1_9_6, 0.0_2_2_5, -0.0_6_9_0, -0.1_7_3_6, 0.0_9_5_1,
						        0.0_9_7_1, -0.0_8_1_7, -0.0_7_0_2, 0.0_1_6_2, 0.0_2_6_0, 0.0_0_1_7, -0.0_1_9_2, -0.1_6_7_8,
						        0.0_7_0_9, -0.1_8_6_7, -0.0_6_5_5, -0.0_2_7_4, -0.0_2_3_4, -0.1_8_8_4, -0.0_5_1_6, -0.0_5_5_4,
						        -0.0_2_7_4, -0.1_4_2_5, -0.1_4_2_3, 0.0_8_3_7, 0.0_3_7_7, -0.0_8_5_4
						    ])
						# fmt: on

						_UpperCAmelCase	 =     self._load_datasamples(1)
						_UpperCAmelCase	 =     WhisperFeatureExtractor()
						_UpperCAmelCase	 =     feature_extractor(A				,	return_tensors='pt').input_features
						self.assertEqual(input_features.shape				,	(1, 80, 30_00))
						self.assertTrue(torch.allclose(input_features[0, 0, :30]				,	A				,	atol=1E-4))



			def 					_lowerCamelCase (		self     :						Dict)     ->					Optional[int]:



						"""simple docstring"""






						_UpperCAmelCase	 =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict())
						_UpperCAmelCase	 =     self._load_datasamples(1)[0]
						_UpperCAmelCase	 =     ((audio - audio.min()) / (audio.max() - audio.min())) * 6_55_35  # Rescale to [0, 65535] to show issue
						_UpperCAmelCase	 =     feat_extract.zero_mean_unit_var_norm([audio]				,	attention_mask=A)[0]

						self.assertTrue(np.all(np.mean(A) < 1E-3))
						self.assertTrue(np.all(np.abs(np.var(A) - 1) < 1E-3))





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 							A			(  _UpperCAmelCase       :     List[Any]					)       ->				int:







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


UpperCAmelCase__     = logging.getLogger(__name__)





def 							A			(  _UpperCAmelCase       :     List[Any]							,     _UpperCAmelCase       :     Union[str, Any]					)       ->				Any:







			'''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=_UpperCAmelCase							,     filename=_UpperCAmelCase							,     monitor=F"val_{metric}"							,     mode='max'							,     save_top_k=1							,     every_n_epochs=1							,     )

			return checkpoint_callback





def 							A			(  _UpperCAmelCase       :     Any							,     _UpperCAmelCase       :     int					)       ->				str:







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







class     __lowerCAmelCase					(    pl.Callback		):



			def 					_lowerCamelCase (		self     :						Optional[int]				,	A     :						List[Any]				,	A     :						int)     ->					Dict:



						"""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 					_lowerCamelCase (		self     :						Optional[Any]				,	A     :						pl.Trainer				,	A     :						pl.LightningModule				,	A     :						str				,	A     :						int=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 					_lowerCamelCase (		self     :						str				,	A     :						Optional[int]				,	A     :						List[str])     ->					Optional[Any]:



						"""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 					_lowerCamelCase (		self     :						Dict				,	A     :						pl.Trainer				,	A     :						pl.LightningModule)     ->					int:



						"""simple docstring"""






						save_json(pl_module.metrics				,	pl_module.metrics_save_path)
						return self._write_logs(A				,	A				,	'test')



			@rank_zero_only
			def 					_lowerCamelCase (		self     :						Tuple				,	A     :						pl.Trainer				,	A     :						str)     ->					Dict:



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