
"""simple docstring"""






import itertools
import random
import unittest

import numpy as np

from transformers import ASTFeatureExtractor
from transformers.testing_utils import require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


UpperCAmelCase        =				random.Random()

if is_torch_available():
					import torch




def __magic_name__	(				_lowerCamelCase:		Optional[int],					_lowerCamelCase:		Tuple=1.0,					_lowerCamelCase:		Optional[int]=None,					_lowerCamelCase:		int=None			)					->						List[Any]:

	'''simple docstring'''




	if rng is None:
		lowerCAmelCase												=  global_rng

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

	return values



class 			lowercase					( unittest.TestCase		):






			def __init__(self  :	Tuple      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	Tuple=7      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]=400      ,SCREAMING_SNAKE_CASE_  :	List[str]=2_000      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]=1      ,SCREAMING_SNAKE_CASE_  :	List[str]=0.0      ,SCREAMING_SNAKE_CASE_  :	Dict=16_000      ,SCREAMING_SNAKE_CASE_  :	int=True      ,SCREAMING_SNAKE_CASE_  :	Any=True      ,)       ->	Any:





				"""simple docstring"""


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






			def 					UpperCAmelCase							(self  :	List[Any]			)       ->	Optional[Any]:





				"""simple docstring"""


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






			def 					UpperCAmelCase							(self  :	str      ,SCREAMING_SNAKE_CASE_  :	List[Any]=False      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]=False			)       ->	int:





				"""simple docstring"""


				def _flatten(SCREAMING_SNAKE_CASE_  :	List[str]			):
					return list(itertools.chain(*SCREAMING_SNAKE_CASE_			)			)

				if equal_length:
					lowerCAmelCase												=  floats_list((self.batch_size, self.max_seq_length)			)
				else:
					# make sure that inputs increase in size
					lowerCAmelCase												=  [
					    _flatten(floats_list((x, self.feature_size)			)			)
					    for x in range(self.min_seq_length      ,self.max_seq_length      ,self.seq_length_diff			)
					]

				if numpify:
					lowerCAmelCase												=  [np.asarray(SCREAMING_SNAKE_CASE_			) for x in speech_inputs]

				return speech_inputs





@require_torch
@require_torchaudio
class 			lowercase					( lowercase__   ,unittest.TestCase		):
			lowercase        =     ASTFeatureExtractor






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





				"""simple docstring"""


				lowerCAmelCase												=  ASTFeatureExtractionTester(self			)






			def 					UpperCAmelCase							(self  :	Any			)       ->	str:





				"""simple docstring"""


				lowerCAmelCase												=  self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()			)
				# create three inputs of length 800, 1000, and 1200
				lowerCAmelCase												=  [floats_list((1, x)			)[0] for x in range(800      ,1_400      ,200			)]
				lowerCAmelCase												=  [np.asarray(SCREAMING_SNAKE_CASE_			) for speech_input in speech_inputs]

				# Test not batched input
				lowerCAmelCase												=  feat_extract(speech_inputs[0]      ,return_tensors='''np'''			).input_values
				lowerCAmelCase												=  feat_extract(np_speech_inputs[0]      ,return_tensors='''np'''			).input_values
				self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,atol=1e-3			)			)

				# Test batched
				lowerCAmelCase												=  feat_extract(SCREAMING_SNAKE_CASE_      ,padding=SCREAMING_SNAKE_CASE_      ,return_tensors='''np'''			).input_values
				lowerCAmelCase												=  feat_extract(SCREAMING_SNAKE_CASE_      ,padding=SCREAMING_SNAKE_CASE_      ,return_tensors='''np'''			).input_values
				for enc_seq_a, enc_seq_a in zip(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			):
					self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,atol=1e-3			)			)

				# Test 2-D numpy arrays are batched.
				lowerCAmelCase												=  [floats_list((1, x)			)[0] for x in (800, 800, 800)]
				lowerCAmelCase												=  np.asarray(SCREAMING_SNAKE_CASE_			)
				lowerCAmelCase												=  feat_extract(SCREAMING_SNAKE_CASE_      ,return_tensors='''np'''			).input_values
				lowerCAmelCase												=  feat_extract(SCREAMING_SNAKE_CASE_      ,return_tensors='''np'''			).input_values
				for enc_seq_a, enc_seq_a in zip(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			):
					self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,atol=1e-3			)			)






			@require_torch
			def 					UpperCAmelCase							(self  :	List[str]			)       ->	Optional[Any]:





				"""simple docstring"""


				import torch

				lowerCAmelCase												=  self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()			)
				lowerCAmelCase												=  np.random.rand(100			).astype(np.floataa			)
				lowerCAmelCase												=  np_speech_inputs.tolist()

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






			def 					UpperCAmelCase							(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]			)       ->	Optional[Any]:





				"""simple docstring"""


				from datasets import load_dataset

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

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






			@require_torch
			def 					UpperCAmelCase							(self  :	int			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase												=  torch.tensor(
				    [-0.98_94, -1.27_76, -0.90_66, -1.27_76, -0.93_49, -1.26_09, -1.03_86, -1.27_76,
				     -1.15_61, -1.27_76, -1.20_52, -1.27_23, -1.21_90, -1.21_32, -1.27_76, -1.11_33,
				     -1.19_53, -1.13_43, -1.15_84, -1.22_03, -1.17_70, -1.24_74, -1.23_81, -1.19_36,
				     -0.92_70, -0.83_17, -0.80_49, -0.77_06, -0.75_65, -0.78_69]			)
				# fmt: on

				lowerCAmelCase												=  self._load_datasamples(1			)
				lowerCAmelCase												=  ASTFeatureExtractor()
				lowerCAmelCase												=  feature_extractor(SCREAMING_SNAKE_CASE_      ,return_tensors='''pt'''			).input_values
				self.assertEquals(input_values.shape      ,(1, 1_024, 128)			)
				self.assertTrue(torch.allclose(input_values[0, 0, :30]      ,SCREAMING_SNAKE_CASE_      ,atol=1e-4			)			)




"""simple docstring"""






from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_torch_dist_unique_port,
    require_torch_multi_gpu,
    require_torch_neuroncore,
)
from transformers.training_args import ParallelMode
from transformers.utils import logging


UpperCAmelCase        =				logging.get_logger(__name__)


if is_torch_available():
					import torch
					from torch import nn
					from torch.utils.data import Dataset

					from transformers import Trainer



					class 			lowercase					( lowercase__		):






								def __init__(self  :	int      ,SCREAMING_SNAKE_CASE_  :	int = 101			)       ->	List[str]:





									"""simple docstring"""


									lowerCAmelCase												=  length






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





									"""simple docstring"""


									return self.length






								def __getitem__(self  :	str      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]			)       ->	int:





									"""simple docstring"""


									return i





					class 			lowercase					:






								def __call__(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	List[Any]			)       ->	Optional[Any]:





									"""simple docstring"""


									return {"input_ids": torch.tensor(SCREAMING_SNAKE_CASE_			), "labels": torch.tensor(SCREAMING_SNAKE_CASE_			)}





					class 			lowercase					( nn.Module		):






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





									"""simple docstring"""


									super().__init__()
									# Add some (unused) params otherwise DDP will complain.
									lowerCAmelCase												=  nn.Linear(120      ,80			)






								def 					UpperCAmelCase							(self  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=None			)       ->	Union[str, Any]:





									"""simple docstring"""


									if labels is not None:
										return torch.tensor(0.0      ,device=input_ids.device			), input_ids
									else:
										return input_ids



class 			lowercase					( lowercase__		):






			@require_torch_neuroncore
			def 					UpperCAmelCase							(self  :	Union[str, Any]			)       ->	Union[str, Any]:





				"""simple docstring"""


				lowerCAmelCase												=  F"""--nproc_per_node=2
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
				lowerCAmelCase												=  self.get_auto_remove_tmp_dir()
				lowerCAmelCase												=  F"""--output_dir {output_dir}""".split()
				lowerCAmelCase												=  ['''torchrun'''] + distributed_args + args
				execute_subprocess_async(SCREAMING_SNAKE_CASE_      ,env=self.get_env()			)
				# successful return here == success - any errors would have caused an error in the sub-call





class 			lowercase					( lowercase__		):






			@require_torch_multi_gpu
			def 					UpperCAmelCase							(self  :	List[str]			)       ->	List[str]:





				"""simple docstring"""


				lowerCAmelCase												=  F"""--nproc_per_node={torch.cuda.device_count()}
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
				lowerCAmelCase												=  self.get_auto_remove_tmp_dir()
				lowerCAmelCase												=  F"""--output_dir {output_dir}""".split()
				lowerCAmelCase												=  ['''torchrun'''] + distributed_args + args
				execute_subprocess_async(SCREAMING_SNAKE_CASE_      ,env=self.get_env()			)
				# successful return here == success - any errors would have caused an error in the sub-call


if __name__ == "__main__":
					# The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:
					#
					# PYTHONPATH="src" python -m torch.distributed.run --nproc_per_node 2 --output_dir output_dir ./tests/test_trainer_distributed.py

					UpperCAmelCase        =				HfArgumentParser((TrainingArguments,))
					UpperCAmelCase        =				parser.parse_args_into_dataclasses()[0]

					logger.warning(
					    f"""Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, """
					    f"""distributed training: {training_args.parallel_mode != ParallelMode.NOT_DISTRIBUTED}"""
					)

					# Essentially, what we want to verify in the distributed case is that we get all samples back,
					# in the right order. (this is crucial for prediction for instance)
					for dataset_length in [1_0_1, 4_0, 7]:
										UpperCAmelCase        =				DummyDataset(dataset_length)




										def __magic_name__	(				_lowerCamelCase:		EvalPrediction			)					->						Dict:

											'''simple docstring'''




											lowerCAmelCase												=  list(range(len(_lowerCamelCase			)			)			)
											lowerCAmelCase												=  p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
											if not success and training_args.local_rank == 0:
												logger.warning(
												    '''Predictions and/or labels do not match expected results:\n  - predictions: '''
												    F"""{p.predictions.tolist()}\n  - labels: {p.label_ids.tolist()}\n  - expected: {sequential}"""			)
											return {"success": success}

										UpperCAmelCase        =				Trainer(
										    model=DummyModel(),
										    args=training_args,
										    data_collator=DummyDataCollator(),
										    eval_dataset=dataset,
										    compute_metrics=compute_metrics,
										)
										UpperCAmelCase        =				trainer.evaluate()
										logger.info(metrics)
										if metrics["eval_success"] is not True:
															logger.error(metrics)
															exit(1)

										UpperCAmelCase        =				trainer.predict(dataset)
										logger.info(p.metrics)
										if p.metrics["test_success"] is not True:
															logger.error(p.metrics)
															exit(1)

										UpperCAmelCase        =				2

										UpperCAmelCase        =				trainer.evaluate()
										logger.info(metrics)
										if metrics["eval_success"] is not True:
															logger.error(metrics)
															exit(1)

										UpperCAmelCase        =				trainer.predict(dataset)
										logger.info(p.metrics)
										if p.metrics["test_success"] is not True:
															logger.error(p.metrics)
															exit(1)



										UpperCAmelCase        =				None


