



"""simple docstring"""







import itertools
import random
import unittest

import numpy as np

from transformers import is_speech_available
from transformers.testing_utils import require_torch, require_torchaudio

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
			from transformers import SpeechaTextFeatureExtractor

UpperCamelCase     =							random.Random()

def       _lowerCamelCase				(							UpperCAmelCase_ :       Any,	UpperCAmelCase_ :       int=1.0,	UpperCAmelCase_ :       Any=None,	UpperCAmelCase_ :       List[Any]=None     )	-> Dict:






	"""simple docstring"""





	if rng is None:
		A__											=					global_rng

	A__											=					[]
	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       UpperCamelCase__			(   unittest.TestCase      ):

			"""simple docstring"""


			def __init__(      self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__=7						,	SCREAMING_SNAKE_CASE__=400						,	SCREAMING_SNAKE_CASE__=2000						,	SCREAMING_SNAKE_CASE__=24						,	SCREAMING_SNAKE_CASE__=24						,	SCREAMING_SNAKE_CASE__=0.0						,	SCREAMING_SNAKE_CASE__=16000						,	SCREAMING_SNAKE_CASE__=True						,	SCREAMING_SNAKE_CASE__=True						,	)      ->	Optional[int]:
				A__											=					parent
				A__											=					batch_size
				A__											=					min_seq_length
				A__											=					max_seq_length
				A__											=					(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
				A__											=					feature_size
				A__											=					num_mel_bins
				A__											=					padding_value
				A__											=					sampling_rate
				A__											=					return_attention_mask
				A__											=					do_normalize


			def   snake_case__	(      self		)      ->	Optional[int]:
				return {
				    "feature_size": self.feature_size,
				    "num_mel_bins": self.num_mel_bins,
				    "padding_value": self.padding_value,
				    "sampling_rate": self.sampling_rate,
				    "return_attention_mask": self.return_attention_mask,
				    "do_normalize": self.do_normalize,
				}







			def   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__=False						,	SCREAMING_SNAKE_CASE__=False		)      ->	Tuple:
				def _flatten(SCREAMING_SNAKE_CASE__		):
					return list(itertools.chain(*SCREAMING_SNAKE_CASE__		)		)

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






@require_torch
@require_torchaudio
class       UpperCamelCase__			(   _lowerCAmelCase       ,		unittest.TestCase      ):

			"""simple docstring"""


			A__	:					str							  =     SpeechaTextFeatureExtractor if is_speech_available() else None


			def   snake_case__	(      self		)      ->	str:
				A__											=					SpeechaTextFeatureExtractionTester(self		)


			def   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__		)      ->	Dict:
				self.assertTrue(np.all(np.mean(SCREAMING_SNAKE_CASE__						,	axis=0		) < 1e-3		)		)
				self.assertTrue(np.all(np.abs(np.var(SCREAMING_SNAKE_CASE__						,	axis=0		) - 1		) < 1e-3		)		)


			def   snake_case__	(      self		)      ->	List[Any]:
				# Tests that all call wrap to encode_plus and batch_encode_plus
				A__											=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
				# create three inputs of length 800, 1000, and 1200
				A__											=					[floats_list((1, x)		)[0] for x in range(800						,	1400						,	200		)]
				A__											=					[np.asarray(SCREAMING_SNAKE_CASE__		) for speech_input in speech_inputs]

				# Test feature size
				A__											=					feature_extractor(SCREAMING_SNAKE_CASE__						,	padding=SCREAMING_SNAKE_CASE__						,	return_tensors="np"		).input_features
				self.assertTrue(input_features.ndim == 3		)
				self.assertTrue(input_features.shape[-1] == feature_extractor.feature_size		)

				# Test not batched input
				A__											=					feature_extractor(speech_inputs[0]						,	return_tensors="np"		).input_features
				A__											=					feature_extractor(np_speech_inputs[0]						,	return_tensors="np"		).input_features
				self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	atol=1e-3		)		)

				# Test batched
				A__											=					feature_extractor(SCREAMING_SNAKE_CASE__						,	return_tensors="np"		).input_features
				A__											=					feature_extractor(SCREAMING_SNAKE_CASE__						,	return_tensors="np"		).input_features
				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.
				A__											=					[floats_list((1, x)		)[0] for x in (800, 800, 800)]
				A__											=					np.asarray(SCREAMING_SNAKE_CASE__		)
				A__											=					feature_extractor(SCREAMING_SNAKE_CASE__						,	return_tensors="np"		).input_features
				A__											=					feature_extractor(SCREAMING_SNAKE_CASE__						,	return_tensors="np"		).input_features
				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		)		)


			def   snake_case__	(      self		)      ->	Any:
				A__											=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
				A__											=					[floats_list((1, x)		)[0] for x in range(800						,	1400						,	200		)]

				A__											=					["longest", "max_length", "do_not_pad"]
				A__											=					[None, 16, None]
				for max_length, padding in zip(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__		):
					A__											=					feature_extractor(
					    SCREAMING_SNAKE_CASE__						,	padding=SCREAMING_SNAKE_CASE__						,	max_length=SCREAMING_SNAKE_CASE__						,	return_attention_mask=SCREAMING_SNAKE_CASE__		)
					A__											=					inputs.input_features
					A__											=					inputs.attention_mask
					A__											=					[np.sum(SCREAMING_SNAKE_CASE__		) for x in attention_mask]

					self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]		)
					self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]		)
					self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]		)


			def   snake_case__	(      self		)      ->	List[Any]:
				A__											=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
				A__											=					[floats_list((1, x)		)[0] for x in range(800						,	1400						,	200		)]

				A__											=					["longest", "max_length", "do_not_pad"]
				A__											=					[None, 16, None]
				for max_length, padding in zip(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__		):
					A__											=					feature_extractor(
					    SCREAMING_SNAKE_CASE__						,	max_length=SCREAMING_SNAKE_CASE__						,	padding=SCREAMING_SNAKE_CASE__						,	return_tensors="np"						,	return_attention_mask=SCREAMING_SNAKE_CASE__		)
					A__											=					inputs.input_features
					A__											=					inputs.attention_mask
					A__											=					[np.sum(SCREAMING_SNAKE_CASE__		) for x in attention_mask]

					self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]		)
					self.assertTrue(input_features[0][fbank_feat_lengths[0] :].sum() < 1e-6		)
					self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]		)
					self.assertTrue(input_features[0][fbank_feat_lengths[1] :].sum() < 1e-6		)
					self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]		)


			def   snake_case__	(      self		)      ->	List[Any]:
				A__											=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
				A__											=					[floats_list((1, x)		)[0] for x in range(800						,	1400						,	200		)]
				A__											=					feature_extractor(
				    SCREAMING_SNAKE_CASE__						,	padding="max_length"						,	max_length=4						,	truncation=SCREAMING_SNAKE_CASE__						,	return_tensors="np"						,	return_attention_mask=SCREAMING_SNAKE_CASE__						,	)
				A__											=					inputs.input_features
				A__											=					inputs.attention_mask
				A__											=					np.sum(attention_mask == 1						,	axis=1		)

				self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]		)
				self._check_zero_mean_unit_variance(input_features[1]		)
				self._check_zero_mean_unit_variance(input_features[2]		)


			def   snake_case__	(      self		)      ->	int:
				A__											=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
				A__											=					[floats_list((1, x)		)[0] for x in range(800						,	1400						,	200		)]
				A__											=					feature_extractor(
				    SCREAMING_SNAKE_CASE__						,	padding="longest"						,	max_length=4						,	truncation=SCREAMING_SNAKE_CASE__						,	return_tensors="np"						,	return_attention_mask=SCREAMING_SNAKE_CASE__						,	)
				A__											=					inputs.input_features
				A__											=					inputs.attention_mask
				A__											=					np.sum(attention_mask == 1						,	axis=1		)

				self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]		)
				self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]		)
				self._check_zero_mean_unit_variance(input_features[2]		)

				# make sure that if max_length < longest -> then pad to max_length
				self.assertEqual(input_features.shape						,	(3, 4, 24)		)

				A__											=					[floats_list((1, x)		)[0] for x in range(800						,	1400						,	200		)]
				A__											=					feature_extractor(
				    SCREAMING_SNAKE_CASE__						,	padding="longest"						,	max_length=16						,	truncation=SCREAMING_SNAKE_CASE__						,	return_tensors="np"						,	return_attention_mask=SCREAMING_SNAKE_CASE__						,	)
				A__											=					inputs.input_features
				A__											=					inputs.attention_mask
				A__											=					np.sum(attention_mask == 1						,	axis=1		)

				self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]		)
				self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]		)
				self._check_zero_mean_unit_variance(input_features[2]		)

				# make sure that if max_length < longest -> then pad to max_length
				self.assertEqual(input_features.shape						,	(3, 6, 24)		)


			def   snake_case__	(      self		)      ->	Optional[int]:
				import torch

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

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


			def   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__		)      ->	List[str]:
				from datasets import load_dataset

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

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







			def   snake_case__	(      self		)      ->	Union[str, Any]:
				# fmt: off
				A__											=					np.array([
				    -1.5_7_4_5, -1.7_7_1_3, -1.7_0_2_0, -1.6_0_6_9, -1.2_2_5_0, -1.1_1_0_5, -0.9_0_7_2, -0.8_2_4_1,
				    -1.2_3_1_0, -0.8_0_9_8, -0.3_3_2_0, -0.4_1_0_1, -0.7_9_8_5, -0.4_9_9_6, -0.8_2_1_3, -0.9_1_2_8,
				    -1.0_4_2_0, -1.1_2_8_6, -1.0_4_4_0, -0.7_9_9_9, -0.8_4_0_5, -1.2_2_7_5, -1.5_4_4_3, -1.4_6_2_5,
				]		)
				# fmt: on

				A__											=					self._load_datasamples(1		)
				A__											=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
				A__											=					feature_extractor(SCREAMING_SNAKE_CASE__						,	return_tensors="pt"		).input_features
				self.assertEquals(input_features.shape						,	(1, 584, 24)		)
				self.assertTrue(np.allclose(input_features[0, 0, :30]						,	SCREAMING_SNAKE_CASE__						,	atol=1e-4		)		)











# XXX: we want transformers master here - in the absense of conftest manipulating sys.path:
# hack it in for now:
import sys
from pathlib import Path


SCREAMING_SNAKE_CASE__							:      List[str]         =    Path(__file__).resolve().parents[3] / """src"""
sys.path.insert(1, str(git_repo_path))

import dataclasses  # noqa
import io  # noqa
import itertools  # noqa
import json  # noqa
import os  # noqa
import unittest  # noqa
from copy import deepcopy  # noqa

from parameterized import parameterized  # noqa
from transformers import TrainingArguments, is_torch_available  # noqa
from transformers.deepspeed import is_deepspeed_available  # noqa
from transformers.file_utils import WEIGHTS_NAME  # noqa
from transformers.testing_utils import (  # noqa
    CaptureLogger,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    mockenv_context,
    require_deepspeed,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.trainer_utils import set_seed  # noqa


set_seed(4_2)

SCREAMING_SNAKE_CASE__							:      Optional[Any]         =    {"""base""": """patrickvonplaten/wav2vec2_tiny_random""", """robust""": """patrickvonplaten/wav2vec2_tiny_random_robust"""}

SCREAMING_SNAKE_CASE__							:      Any         =    """zero2"""
SCREAMING_SNAKE_CASE__							:      Dict         =    """zero3"""
SCREAMING_SNAKE_CASE__							:      str         =    [ZEROa, ZEROa]






def    _A	(   lowerCamelCase				,			lowerCamelCase				,			lowerCamelCase ):
   # customize the test name generator function as we want both params to appear in the sub-test
   # name, as by default it shows only the first param
   a__						:    Optional[Any]   		=	parameterized.to_safe_name("_".join(str(lowerCamelCase ) for x in param.args ) )
   return F"""{func.__name__}_{param_based_name}"""


# Cartesian-product of zero stages with models to test
SCREAMING_SNAKE_CASE__							:      Union[str, Any]         =    list(itertools.product(stages, models.keys()))







@slow
@require_deepspeed
@require_torch_gpu
class 	__lowerCAmelCase      (       _UpperCamelCase	):





   @parameterized.expand(snake_case       ,      name_func=snake_case   )
   def 						_snake_case			(			self       ,      snake_case       ,      snake_case   )					->	Dict:
      """simple docstring"""





      self.run_and_check(
          stage=snake_case       ,      model=snake_case       ,      distributed=snake_case       ,      fpaa=snake_case       ,      )





   @require_torch_multi_gpu
   @parameterized.expand(snake_case       ,      name_func=snake_case   )
   def 						_snake_case			(			self       ,      snake_case       ,      snake_case   )					->	str:
      """simple docstring"""





      self.run_and_check(
          stage=snake_case       ,      model=snake_case       ,      distributed=snake_case       ,      fpaa=snake_case       ,      )





   @parameterized.expand(snake_case       ,      name_func=snake_case   )
   def 						_snake_case			(			self       ,      snake_case       ,      snake_case   )					->	List[str]:
      """simple docstring"""





      self.run_and_check(
          stage=snake_case       ,      model=snake_case       ,      distributed=snake_case       ,      fpaa=snake_case       ,      )





   @require_torch_multi_gpu
   @parameterized.expand(snake_case       ,      name_func=snake_case   )
   def 						_snake_case			(			self       ,      snake_case       ,      snake_case   )					->	List[str]:
      """simple docstring"""





      self.run_and_check(
          stage=snake_case       ,      model=snake_case       ,      distributed=snake_case       ,      fpaa=snake_case       ,      )





   def 						_snake_case			(			self       ,      snake_case   )					->	str:
      """simple docstring"""





      pass





   def 						_snake_case			(			self       ,      snake_case       ,      snake_case       ,      snake_case = 10       ,      snake_case = True       ,      snake_case = True       ,      snake_case = True       ,      )					->	str:
      """simple docstring"""





      a__						:    Tuple   		=	models[model]

      a__						:    int   		=	self.run_trainer(
          stage=snake_case       ,      model_name=snake_case       ,      eval_steps=snake_case       ,      num_train_epochs=1       ,      distributed=snake_case       ,      fpaa=snake_case       ,      )

      self.do_checks(snake_case   )

      return output_dir





   def 						_snake_case			(			self       ,      snake_case       ,      snake_case       ,      snake_case = 10       ,      snake_case = 1       ,      snake_case = True       ,      snake_case = True       ,      )					->	Optional[Any]:
      """simple docstring"""





      a__						:    str   		=	self.get_auto_remove_tmp_dir("./xxx"       ,      after=snake_case   )
      a__						:    List[Any]   		=	F"""
            --model_name_or_path {model_name}
            --dataset_name hf-internal-testing/librispeech_asr_dummy
            --dataset_config_name clean
            --train_split_name validation
            --validation_split_name validation
            --output_dir {output_dir}
            --num_train_epochs {str(snake_case   )}
            --per_device_train_batch_size 2
            --per_device_eval_batch_size 2
            --evaluation_strategy steps
            --learning_rate 5e-4
            --warmup_steps 8
            --orthography timit
            --preprocessing_num_workers 1
            --group_by_length
            --freeze_feature_extractor
            --report_to none
            --save_steps 0
            --eval_steps {eval_steps}
            --report_to none
        """.split()

      if fpaa:
         args.extend(["--fp16"]   )

      # currently ds_config_wav2vec2_zero.json requires "zero_optimization.find_unused_parameters": true,
      # hence the separate config files
      a__						:    str   		=	F"""--deepspeed {self.test_file_dir_str}/ds_config_wav2vec2_{stage}.json""".split()
      a__						:    Dict   		=	[F"""{self.examples_dir_str}/research_projects/wav2vec2/run_asr.py"""]
      a__						:    Optional[int]   		=	self.get_launcher(snake_case   )

      a__						:    Optional[int]   		=	launcher + script + args + ds_args
      # keep for quick debug
      # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
      execute_subprocess_async(snake_case       ,      env=self.get_env()   )

      return output_dir





   def 						_snake_case			(			self       ,      snake_case=False   )					->	List[Any]:
      """simple docstring"""





      a__						:    Union[str, Any]   		=	min(2       ,      get_gpu_count()   ) if distributed else 1
      return F"""deepspeed --num_nodes 1 --num_gpus {num_gpus}""".split()

