'''simple docstring'''

import itertools
import random
import unittest

import numpy as np

from transformers import BatchFeature, SpeechTaFeatureExtractor
from transformers.testing_utils import require_torch
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_torch_available():
  import torch


UpperCamelCase__  :						int					=      random.Random()







def __UpperCamelCase( _A       :					Any	,      _A       :					Union[str, Any]=1.0	,      _A       :					str=None	,      _A       :					List[str]=None     ):
   '''simple docstring'''
   if rng is None:
      UpperCAmelCase__				:					int 					=	global_rng

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

   return values


@require_torch
class 					_lowercase						(   unittest.TestCase ):
      '''simple docstring'''


      def __init__(					self       ,lowerCamelCase_       ,lowerCamelCase_=7       ,lowerCamelCase_=400       ,lowerCamelCase_=2000       ,lowerCamelCase_=1       ,lowerCamelCase_=0.0       ,lowerCamelCase_=16000       ,lowerCamelCase_=True       ,lowerCamelCase_=80       ,lowerCamelCase_=16       ,lowerCamelCase_=64       ,lowerCamelCase_="hann_window"       ,lowerCamelCase_=80       ,lowerCamelCase_=7600       ,lowerCamelCase_=1e-10       ,lowerCamelCase_=True       ,)      ->      Union[str, Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					List[str] 					=	parent
         UpperCAmelCase__				:					Tuple 					=	batch_size
         UpperCAmelCase__				:					Optional[int] 					=	min_seq_length
         UpperCAmelCase__				:					int 					=	max_seq_length
         UpperCAmelCase__				:					List[Any] 					=	(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
         UpperCAmelCase__				:					List[Any] 					=	feature_size
         UpperCAmelCase__				:					List[str] 					=	padding_value
         UpperCAmelCase__				:					Optional[Any] 					=	sampling_rate
         UpperCAmelCase__				:					Optional[int] 					=	do_normalize
         UpperCAmelCase__				:					Tuple 					=	num_mel_bins
         UpperCAmelCase__				:					Optional[Any] 					=	hop_length
         UpperCAmelCase__				:					Any 					=	win_length
         UpperCAmelCase__				:					Optional[int] 					=	win_function
         UpperCAmelCase__				:					List[str] 					=	fmin
         UpperCAmelCase__				:					str 					=	fmax
         UpperCAmelCase__				:					Optional[Any] 					=	mel_floor
         UpperCAmelCase__				:					Any 					=	return_attention_mask




      def 		lowerCAmelCase__       (					self				)      ->      Dict:



         '''simple docstring'''



         return {
             "feature_size": self.feature_size,
             "padding_value": self.padding_value,
             "sampling_rate": self.sampling_rate,
             "do_normalize": self.do_normalize,
             "num_mel_bins": self.num_mel_bins,
             "hop_length": self.hop_length,
             "win_length": self.win_length,
             "win_function": self.win_function,
             "fmin": self.fmin,
             "fmax": self.fmax,
             "mel_floor": self.mel_floor,
             "return_attention_mask": self.return_attention_mask,
         }




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_=False       ,lowerCamelCase_=False				)      ->      Tuple:



         '''simple docstring'''



         def _flatten(lowerCamelCase_				):
            return list(itertools.chain(*lowerCamelCase_				)				)

         if equal_length:
            UpperCAmelCase__				:					Any 					=	floats_list((self.batch_size, self.max_seq_length)				)
         else:
            # make sure that inputs increase in size
            UpperCAmelCase__				:					Optional[Any] 					=	[
                _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:
            UpperCAmelCase__				:					List[Any] 					=	[np.asarray(lowerCamelCase_				) for x in speech_inputs]

         return speech_inputs




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_=False       ,lowerCamelCase_=False				)      ->      List[str]:



         '''simple docstring'''



         if equal_length:
            UpperCAmelCase__				:					Optional[int] 					=	[floats_list((self.max_seq_length, self.num_mel_bins)				) for _ in range(self.batch_size				)]
         else:
            # make sure that inputs increase in size
            UpperCAmelCase__				:					Tuple 					=	[
                floats_list((x, self.num_mel_bins)				)
                for x in range(self.min_seq_length       ,self.max_seq_length       ,self.seq_length_diff				)
            ]

         if numpify:
            UpperCAmelCase__				:					Tuple 					=	[np.asarray(lowerCamelCase_				) for x in speech_inputs]

         return speech_inputs




@require_torch
class 					_lowercase						(   lowerCAmelCase		,unittest.TestCase ):
      '''simple docstring'''


      UpperCAmelCase_   :   List[Any]   				= SpeechTaFeatureExtractor




      def 		lowerCAmelCase__       (					self				)      ->      str:



         '''simple docstring'''



         UpperCAmelCase__				:					List[str] 					=	SpeechTaFeatureExtractionTester(self				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_				)      ->      str:



         '''simple docstring'''



         self.assertTrue(np.all(np.mean(lowerCamelCase_       ,axis=0				) < 1e-3				)				)
         self.assertTrue(np.all(np.abs(np.var(lowerCamelCase_       ,axis=0				) - 1				) < 1e-3				)				)




      def 		lowerCAmelCase__       (					self				)      ->      Tuple:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[int] 					=	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()				)
         # create three inputs of length 800, 1000, and 1200
         UpperCAmelCase__				:					Dict 					=	[floats_list((1, x)				)[0] for x in range(800       ,1400       ,200				)]
         UpperCAmelCase__				:					int 					=	[np.asarray(lowerCamelCase_				) for speech_input in speech_inputs]

         # Test not batched input
         UpperCAmelCase__				:					Optional[Any] 					=	feat_extract(speech_inputs[0]       ,return_tensors='''np'''				).input_values
         UpperCAmelCase__				:					List[Any] 					=	feat_extract(np_speech_inputs[0]       ,return_tensors='''np'''				).input_values
         self.assertTrue(np.allclose(lowerCamelCase_       ,lowerCamelCase_       ,atol=1e-3				)				)

         # Test batched
         UpperCAmelCase__				:					Union[str, Any] 					=	feat_extract(lowerCamelCase_       ,return_tensors='''np'''				).input_values
         UpperCAmelCase__				:					List[str] 					=	feat_extract(lowerCamelCase_       ,return_tensors='''np'''				).input_values
         for enc_seq_a, enc_seq_a in zip(lowerCamelCase_       ,lowerCamelCase_				):
            self.assertTrue(np.allclose(lowerCamelCase_       ,lowerCamelCase_       ,atol=1e-3				)				)




      def 		lowerCAmelCase__       (					self				)      ->      int:



         '''simple docstring'''



         UpperCAmelCase__				:					List[str] 					=	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()				)
         UpperCAmelCase__				:					Optional[Any] 					=	[floats_list((1, x)				)[0] for x in range(800       ,1400       ,200				)]

         UpperCAmelCase__				:					Union[str, Any] 					=	['''longest''', '''max_length''', '''do_not_pad''']
         UpperCAmelCase__				:					Optional[Any] 					=	[None, 1600, None]
         for max_length, padding in zip(lowerCamelCase_       ,lowerCamelCase_				):
            UpperCAmelCase__				:					str 					=	feat_extract(lowerCamelCase_       ,padding=lowerCamelCase_       ,max_length=lowerCamelCase_       ,return_tensors='''np'''				)
            UpperCAmelCase__				:					Any 					=	processed.input_values

            self._check_zero_mean_unit_variance(input_values[0][:800]				)
            self.assertTrue(input_values[0][800:].sum() < 1e-6				)
            self._check_zero_mean_unit_variance(input_values[1][:1000]				)
            self.assertTrue(input_values[0][1000:].sum() < 1e-6				)
            self._check_zero_mean_unit_variance(input_values[2][:1200]				)




      def 		lowerCAmelCase__       (					self				)      ->      Any:



         '''simple docstring'''



         UpperCAmelCase__				:					str 					=	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()				)
         UpperCAmelCase__				:					Optional[Any] 					=	range(800       ,1400       ,200				)
         UpperCAmelCase__				:					Dict 					=	[floats_list((1, x)				)[0] for x in lengths]

         UpperCAmelCase__				:					List[str] 					=	['''longest''', '''max_length''', '''do_not_pad''']
         UpperCAmelCase__				:					List[str] 					=	[None, 1600, None]

         for max_length, padding in zip(lowerCamelCase_       ,lowerCamelCase_				):
            UpperCAmelCase__				:					Union[str, Any] 					=	feat_extract(lowerCamelCase_       ,max_length=lowerCamelCase_       ,padding=lowerCamelCase_				)
            UpperCAmelCase__				:					Dict 					=	processed.input_values

            self._check_zero_mean_unit_variance(input_values[0][:800]				)
            self._check_zero_mean_unit_variance(input_values[1][:1000]				)
            self._check_zero_mean_unit_variance(input_values[2][:1200]				)




      def 		lowerCAmelCase__       (					self				)      ->      List[str]:



         '''simple docstring'''



         UpperCAmelCase__				:					Union[str, Any] 					=	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()				)
         UpperCAmelCase__				:					Tuple 					=	[floats_list((1, x)				)[0] for x in range(800       ,1400       ,200				)]
         UpperCAmelCase__				:					List[str] 					=	feat_extract(
             lowerCamelCase_       ,truncation=lowerCamelCase_       ,max_length=1000       ,padding='''max_length'''       ,return_tensors='''np'''				)
         UpperCAmelCase__				:					Tuple 					=	processed.input_values

         self._check_zero_mean_unit_variance(input_values[0, :800]				)
         self._check_zero_mean_unit_variance(input_values[1]				)
         self._check_zero_mean_unit_variance(input_values[2]				)




      def 		lowerCAmelCase__       (					self				)      ->      Optional[int]:



         '''simple docstring'''



         UpperCAmelCase__				:					Tuple 					=	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()				)
         UpperCAmelCase__				:					Optional[Any] 					=	[floats_list((1, x)				)[0] for x in range(800       ,1400       ,200				)]
         UpperCAmelCase__				:					Tuple 					=	feat_extract(
             lowerCamelCase_       ,truncation=lowerCamelCase_       ,max_length=1000       ,padding='''longest'''       ,return_tensors='''np'''				)
         UpperCAmelCase__				:					Tuple 					=	processed.input_values

         self._check_zero_mean_unit_variance(input_values[0, :800]				)
         self._check_zero_mean_unit_variance(input_values[1, :1000]				)
         self._check_zero_mean_unit_variance(input_values[2]				)

         # make sure that if max_length < longest -> then pad to max_length
         self.assertTrue(input_values.shape == (3, 1000)				)

         UpperCAmelCase__				:					List[Any] 					=	[floats_list((1, x)				)[0] for x in range(800       ,1400       ,200				)]
         UpperCAmelCase__				:					Optional[Any] 					=	feat_extract(
             lowerCamelCase_       ,truncation=lowerCamelCase_       ,max_length=2000       ,padding='''longest'''       ,return_tensors='''np'''				)
         UpperCAmelCase__				:					Optional[Any] 					=	processed.input_values

         self._check_zero_mean_unit_variance(input_values[0, :800]				)
         self._check_zero_mean_unit_variance(input_values[1, :1000]				)
         self._check_zero_mean_unit_variance(input_values[2]				)

         # make sure that if max_length > longest -> then pad to longest
         self.assertTrue(input_values.shape == (3, 1200)				)




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



         '''simple docstring'''



         UpperCAmelCase__				:					List[Any] 					=	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()				)
         UpperCAmelCase__				:					str 					=	np.random.rand(100				).astype(np.floataa				)
         UpperCAmelCase__				:					Tuple 					=	np_speech_inputs.tolist()

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




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



         '''simple docstring'''



         UpperCAmelCase__				:					str 					=	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()				)
         # create three inputs of length 800, 1000, and 1200
         UpperCAmelCase__				:					List[str] 					=	[floats_list((1, x)				)[0] for x in range(800       ,1400       ,200				)]
         UpperCAmelCase__				:					Dict 					=	[np.asarray(lowerCamelCase_				) for speech_input in speech_inputs]

         # Test feature size
         UpperCAmelCase__				:					int 					=	feature_extractor(audio_target=lowerCamelCase_       ,padding=lowerCamelCase_       ,return_tensors='''np'''				).input_values
         self.assertTrue(input_values.ndim == 3				)
         self.assertTrue(input_values.shape[-1] == feature_extractor.num_mel_bins				)

         # Test not batched input
         UpperCAmelCase__				:					Any 					=	feature_extractor(speech_inputs[0]       ,return_tensors='''np'''				).input_values
         UpperCAmelCase__				:					List[Any] 					=	feature_extractor(np_speech_inputs[0]       ,return_tensors='''np'''				).input_values
         self.assertTrue(np.allclose(lowerCamelCase_       ,lowerCamelCase_       ,atol=1e-3				)				)

         # Test batched
         UpperCAmelCase__				:					Optional[int] 					=	feature_extractor(lowerCamelCase_       ,return_tensors='''np'''				).input_values
         UpperCAmelCase__				:					Optional[int] 					=	feature_extractor(lowerCamelCase_       ,return_tensors='''np'''				).input_values
         for enc_seq_a, enc_seq_a in zip(lowerCamelCase_       ,lowerCamelCase_				):
            self.assertTrue(np.allclose(lowerCamelCase_       ,lowerCamelCase_       ,atol=1e-3				)				)

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




      def 		lowerCAmelCase__       (					self				)      ->      Optional[int]:



         '''simple docstring'''



         UpperCAmelCase__				:					Any 					=	self.feat_extract_tester.prepare_inputs_for_target()
         UpperCAmelCase__				:					List[Any] 					=	self.feature_extraction_class(**self.feat_extract_dict				)
         UpperCAmelCase__				:					Any 					=	feat_extract.model_input_names[0]

         UpperCAmelCase__				:					Dict 					=	BatchFeature({input_name: speech_inputs}				)

         self.assertTrue(all(len(lowerCamelCase_				) == len(lowerCamelCase_				) for x, y in zip(lowerCamelCase_       ,processed_features[input_name]				)				)				)

         UpperCAmelCase__				:					str 					=	self.feat_extract_tester.prepare_inputs_for_target(equal_length=lowerCamelCase_				)
         UpperCAmelCase__				:					str 					=	BatchFeature({input_name: speech_inputs}       ,tensor_type='''np'''				)

         UpperCAmelCase__				:					Union[str, Any] 					=	processed_features[input_name]

         if len(batch_features_input.shape				) < 3:
            UpperCAmelCase__				:					str 					=	batch_features_input[:, :, None]

         self.assertTrue(
             batch_features_input.shape
             == (self.feat_extract_tester.batch_size, len(speech_inputs[0]				), self.feat_extract_tester.num_mel_bins)				)




      @require_torch
      def 		lowerCAmelCase__       (					self				)      ->      List[str]:



         '''simple docstring'''



         UpperCAmelCase__				:					List[Any] 					=	self.feat_extract_tester.prepare_inputs_for_target(equal_length=lowerCamelCase_				)
         UpperCAmelCase__				:					Tuple 					=	self.feature_extraction_class(**self.feat_extract_dict				)
         UpperCAmelCase__				:					List[Any] 					=	feat_extract.model_input_names[0]

         UpperCAmelCase__				:					Union[str, Any] 					=	BatchFeature({input_name: speech_inputs}       ,tensor_type='''pt'''				)

         UpperCAmelCase__				:					Dict 					=	processed_features[input_name]

         if len(batch_features_input.shape				) < 3:
            UpperCAmelCase__				:					List[str] 					=	batch_features_input[:, :, None]

         self.assertTrue(
             batch_features_input.shape
             == (self.feat_extract_tester.batch_size, len(speech_inputs[0]				), self.feat_extract_tester.num_mel_bins)				)




      @require_torch
      def 		lowerCAmelCase__       (					self				)      ->      str:



         '''simple docstring'''



         UpperCAmelCase__				:					Tuple 					=	self.feature_extraction_class(**self.feat_extract_dict				)
         UpperCAmelCase__				:					Any 					=	self.feat_extract_tester.prepare_inputs_for_target()
         UpperCAmelCase__				:					List[str] 					=	feat_extract.model_input_names[0]

         UpperCAmelCase__				:					Tuple 					=	BatchFeature({input_name: speech_inputs}				)

         UpperCAmelCase__				:					Union[str, Any] 					=	feat_extract.num_mel_bins  # hack!

         UpperCAmelCase__				:					int 					=	feat_extract.pad(lowerCamelCase_       ,padding='''longest'''       ,return_tensors='''np'''				)[input_name]
         UpperCAmelCase__				:					Dict 					=	feat_extract.pad(lowerCamelCase_       ,padding='''longest'''       ,return_tensors='''pt'''				)[input_name]

         self.assertTrue(abs(input_np.astype(np.floataa				).sum() - input_pt.numpy().astype(np.floataa				).sum()				) < 1e-2				)




      def 		lowerCAmelCase__       (					self				)      ->      Tuple:



         '''simple docstring'''



         UpperCAmelCase__				:					str 					=	self.feat_extract_dict
         UpperCAmelCase__				:					List[str] 					=	True
         UpperCAmelCase__				:					List[str] 					=	self.feature_extraction_class(**lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[int] 					=	self.feat_extract_tester.prepare_inputs_for_target()
         UpperCAmelCase__				:					List[Any] 					=	[len(lowerCamelCase_				) for x in speech_inputs]
         UpperCAmelCase__				:					Tuple 					=	feat_extract.model_input_names[0]

         UpperCAmelCase__				:					Optional[Any] 					=	BatchFeature({input_name: speech_inputs}				)

         UpperCAmelCase__				:					Dict 					=	feat_extract.num_mel_bins  # hack!

         UpperCAmelCase__				:					str 					=	feat_extract.pad(lowerCamelCase_       ,padding='''longest'''       ,return_tensors='''np'''				)
         self.assertIn('''attention_mask'''       ,lowerCamelCase_				)
         self.assertListEqual(list(processed.attention_mask.shape				)       ,list(processed[input_name].shape[:2]				)				)
         self.assertListEqual(processed.attention_mask.sum(-1				).tolist()       ,lowerCamelCase_				)




      def 		lowerCAmelCase__       (					self				)      ->      int:



         '''simple docstring'''



         UpperCAmelCase__				:					Tuple 					=	self.feat_extract_dict
         UpperCAmelCase__				:					List[Any] 					=	True
         UpperCAmelCase__				:					Optional[int] 					=	self.feature_extraction_class(**lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[Any] 					=	self.feat_extract_tester.prepare_inputs_for_target()
         UpperCAmelCase__				:					List[str] 					=	[len(lowerCamelCase_				) for x in speech_inputs]
         UpperCAmelCase__				:					Optional[Any] 					=	feat_extract.model_input_names[0]

         UpperCAmelCase__				:					Tuple 					=	BatchFeature({input_name: speech_inputs}				)
         UpperCAmelCase__				:					str 					=	min(lowerCamelCase_				)

         UpperCAmelCase__				:					Tuple 					=	feat_extract.num_mel_bins  # hack!

         UpperCAmelCase__				:					Union[str, Any] 					=	feat_extract.pad(
             lowerCamelCase_       ,padding='''max_length'''       ,max_length=lowerCamelCase_       ,truncation=lowerCamelCase_       ,return_tensors='''np'''				)
         self.assertIn('''attention_mask'''       ,lowerCamelCase_				)
         self.assertListEqual(
             list(processed_pad.attention_mask.shape				)       ,[processed_pad[input_name].shape[0], max_length]				)
         self.assertListEqual(
             processed_pad.attention_mask[:, :max_length].sum(-1				).tolist()       ,[max_length for x in speech_inputs]				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_				)      ->      int:



         '''simple docstring'''



         from datasets import load_dataset

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

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




      def 		lowerCAmelCase__       (					self				)      ->      str:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[Any] 					=	torch.tensor(
             [2.3804e-03, 2.0752e-03, 1.9836e-03, 2.1057e-03, 1.6174e-03,
              3.0518e-04, 9.1553e-05, 3.3569e-04, 9.7656e-04, 1.8311e-03,
              2.0142e-03, 2.1057e-03, 1.7395e-03, 4.5776e-04, -3.9673e-04,
              4.5776e-04, 1.0071e-03, 9.1553e-05, 4.8828e-04, 1.1597e-03,
              7.3242e-04, 9.4604e-04, 1.8005e-03, 1.8311e-03, 8.8501e-04,
              4.2725e-04, 4.8828e-04, 7.3242e-04, 1.0986e-03, 2.1057e-03]				)
         # fmt: on

         UpperCAmelCase__				:					str 					=	self._load_datasamples(1				)
         UpperCAmelCase__				:					Optional[Any] 					=	SpeechTaFeatureExtractor()
         UpperCAmelCase__				:					Union[str, Any] 					=	feature_extractor(lowerCamelCase_       ,return_tensors='''pt'''				).input_values
         self.assertEquals(input_values.shape       ,(1, 93680)				)
         self.assertTrue(torch.allclose(input_values[0, :30]       ,lowerCamelCase_       ,atol=1e-6				)				)




      def 		lowerCAmelCase__       (					self				)      ->      Tuple:



         '''simple docstring'''



         UpperCAmelCase__				:					Dict 					=	torch.tensor(
             [-2.6870, -3.0104, -3.1356, -3.5352, -3.0044, -3.0353, -3.4719, -3.6777,
              -3.1520, -2.9435, -2.6553, -2.8795, -2.9944, -2.5921, -3.0279, -3.0386,
              -3.0864, -3.1291, -3.2353, -2.7444, -2.6831, -2.7287, -3.1761, -3.1571,
              -3.2726, -3.0582, -3.1007, -3.4533, -3.4695, -3.0998]				)
         # fmt: on

         UpperCAmelCase__				:					int 					=	self._load_datasamples(1				)
         UpperCAmelCase__				:					List[str] 					=	SpeechTaFeatureExtractor()
         UpperCAmelCase__				:					str 					=	feature_extractor(audio_target=lowerCamelCase_       ,return_tensors='''pt'''				).input_values
         self.assertEquals(input_values.shape       ,(1, 366, 80)				)
         self.assertTrue(torch.allclose(input_values[0, 0, :30]       ,lowerCamelCase_       ,atol=1e-4				)				)





'''simple docstring'''

import logging
import re

import pytorch_quantization
import pytorch_quantization.nn as quant_nn
import torch
from pytorch_quantization import calib
from pytorch_quantization.tensor_quant import QuantDescriptor


UpperCamelCase__  :						Optional[Any]					=      logging.getLogger(__name__)

UpperCamelCase__  :						Dict					=      50  # max width of layer names
UpperCamelCase__  :						Any					=      70  # max width of quantizer names







def __UpperCamelCase( _A       :					Union[str, Any]     ):
   '''simple docstring'''
   UpperCAmelCase__				:					List[str] 					=	parser.add_argument_group('''quant_trainer arguments'''     )
   group.add_argument('''--wprec'''	,      type=_A	,      default=8	,      help='''weight precision'''     )
   group.add_argument('''--aprec'''	,      type=_A	,      default=8	,      help='''activation precision'''     )
   group.add_argument('''--quant-per-tensor'''	,      action='''store_true'''	,      help='''per tensor weight scaling'''     )
   group.add_argument('''--quant-disable'''	,      action='''store_true'''	,      help='''disable all quantizers'''     )
   group.add_argument('''--quant-disable-embeddings'''	,      action='''store_true'''	,      help='''disable all embeddings quantizers'''     )
   group.add_argument('''--quant-disable-keyword'''	,      type=_A	,      nargs='''+'''	,      help='''disable quantizers by keyword'''     )
   group.add_argument('''--quant-disable-layer-module'''	,      type=_A	,      help='''disable quantizers by keyword under layer.'''     )
   group.add_argument('''--quant-enable-layer-module'''	,      type=_A	,      help='''enable quantizers by keyword under layer'''     )
   group.add_argument('''--calibrator'''	,      default='''max'''	,      help='''which quantization range calibrator to use'''     )
   group.add_argument('''--percentile'''	,      default=_A	,      type=_A	,      help='''percentile for PercentileCalibrator'''     )
   group.add_argument('''--fuse-qkv'''	,      action='''store_true'''	,      help='''use the same scale factor for qkv'''     )
   group.add_argument('''--clip-gelu'''	,      metavar='''N'''	,      type=_A	,      help='''clip gelu output maximum value to N'''     )
   group.add_argument(
       '''--recalibrate-weights'''	,      action='''store_true'''	,      help=(
           '''recalibrate weight amaxes by taking the max of the weights.'''
           ''' amaxes will be computed with the current quantization granularity (axis).'''
       )	,      )







def __UpperCamelCase( _A       :					Tuple     ):
   '''simple docstring'''

   if args.calibrator == "max":
      UpperCAmelCase__				:					str 					=	'''max'''
   elif args.calibrator == "percentile":
      if args.percentile is None:
         raise ValueError('''Specify --percentile when using percentile calibrator'''     )
      UpperCAmelCase__				:					Dict 					=	'''histogram'''
   elif args.calibrator == "mse":
      UpperCAmelCase__				:					Any 					=	'''histogram'''
   else:
      raise ValueError(F'''Invalid calibrator {args.calibrator}'''     )

   UpperCAmelCase__				:					Dict 					=	QuantDescriptor(num_bits=args.aprec	,      calib_method=_A     )
   UpperCAmelCase__				:					str 					=	QuantDescriptor(num_bits=args.wprec	,      axis=(None if args.quant_per_tensor else (0,))     )
   quant_nn.QuantLinear.set_default_quant_desc_input(_A     )
   quant_nn.QuantLinear.set_default_quant_desc_weight(_A     )







def __UpperCamelCase( _A       :					Any	,      _A       :					Any	,      _A       :					Any=False	,      _A       :					Optional[Any]=False     ):
   '''simple docstring'''
   logger.info('''Configuring Model for Quantization'''     )
   logger.info(F'''using quantization package {pytorch_quantization.__file__}'''     )

   if not calib:
      if args.quant_disable_embeddings:
         set_quantizer_by_name(_A	,      ['''embeddings''']	,      which='''weight'''	,      _disabled=_A     )

      if args.quant_disable:
         set_quantizer_by_name(_A	,      ['''''']	,      _disabled=_A     )

      if args.quant_disable_keyword:
         set_quantizer_by_name(_A	,      args.quant_disable_keyword	,      _disabled=_A     )

      if args.quant_disable_layer_module:
         set_quantizer_by_name(_A	,      [R'''layer.\d+.''' + args.quant_disable_layer_module]	,      _disabled=_A     )

      if args.quant_enable_layer_module:
         set_quantizer_by_name(_A	,      [R'''layer.\d+.''' + args.quant_enable_layer_module]	,      _disabled=_A     )

      if args.recalibrate_weights:
         recalibrate_weights(_A     )

      if args.fuse_qkv:
         fuse_qkv(_A	,      _A     )

   if args.clip_gelu:
      clip_gelu(_A	,      args.clip_gelu     )

   # if args.local_rank in [-1, 0] and not calib:
   print_quant_summary(_A     )







def __UpperCamelCase( _A       :					str     ):
   '''simple docstring'''
   logger.info('''Enabling Calibration'''     )
   for name, module in model.named_modules():
      if name.endswith('''_quantizer'''     ):
         if module._calibrator is not None:
            module.disable_quant()
            module.enable_calib()
         else:
            module.disable()
         logger.info(F'''{name:80}: {module}'''     )







def __UpperCamelCase( _A       :					Tuple	,      _A       :					Any     ):
   '''simple docstring'''
   logger.info('''Loading calibrated amax'''     )
   for name, module in model.named_modules():
      if name.endswith('''_quantizer'''     ):
         if module._calibrator is not None:
            if isinstance(module._calibrator	,      calib.MaxCalibrator     ):
               module.load_calib_amax()
            else:
               module.load_calib_amax('''percentile'''	,      percentile=args.percentile     )
            module.enable_quant()
            module.disable_calib()
         else:
            module.enable()
   model.cuda()
   print_quant_summary(_A     )







def __UpperCamelCase( _A       :					Dict	,      _A       :					Optional[int]     ):
   '''simple docstring'''

   def fusea(_A       :					Optional[Any]	,      _A       :					Optional[Any]	,      _A       :					Dict     ):
      for mod in [qq, qk, qv]:
         if not hasattr(_A	,      '''_amax'''     ):
            print('''          WARNING: NO AMAX BUFFER'''     )
            return
      UpperCAmelCase__				:					Dict 					=	qq._amax.detach().item()
      UpperCAmelCase__				:					List[Any] 					=	qk._amax.detach().item()
      UpperCAmelCase__				:					Optional[int] 					=	qv._amax.detach().item()

      UpperCAmelCase__				:					Dict 					=	max(_A	,      _A	,      _A     )
      qq._amax.fill_(_A     )
      qk._amax.fill_(_A     )
      qv._amax.fill_(_A     )
      logger.info(F'''          q={q:5.2f} k={k:5.2f} v={v:5.2f} -> {amax:5.2f}'''     )

   for name, mod in model.named_modules():
      if name.endswith('''.attention.self'''     ):
         logger.info(F'''FUSE_QKV: {name:{name_width}}'''     )
         fusea(mod.matmul_q_input_quantizer	,      mod.matmul_k_input_quantizer	,      mod.matmul_v_input_quantizer     )
         if args.quant_per_tensor:
            fusea(mod.query._weight_quantizer	,      mod.key._weight_quantizer	,      mod.value._weight_quantizer     )







def __UpperCamelCase( _A       :					Dict	,      _A       :					Any     ):
   '''simple docstring'''

   for name, mod in model.named_modules():
      if name.endswith('''.output.dense'''     ) and not name.endswith('''attention.output.dense'''     ):
         UpperCAmelCase__				:					Union[str, Any] 					=	mod._input_quantizer._amax.data.detach().item()
         mod._input_quantizer._amax.data.detach().clamp_(max=_A     )
         UpperCAmelCase__				:					Tuple 					=	mod._input_quantizer._amax.data.detach().item()
         logger.info(F'''CLIP_GELU: {name:{name_width}} amax: {amax_init:5.2f} -> {amax:5.2f}'''     )







def __UpperCamelCase( _A       :					str     ):
   '''simple docstring'''

   for name, mod in model.named_modules():
      if hasattr(_A	,      '''_weight_quantizer'''     ) and mod._weight_quantizer.axis is not None:
         UpperCAmelCase__				:					int 					=	mod.weight.shape[0]
         UpperCAmelCase__				:					Tuple 					=	mod._weight_quantizer._amax.detach()
         UpperCAmelCase__				:					Optional[int] 					=	torch.ones(_A	,      dtype=amax.dtype	,      device=amax.device     ) * amax
         print(F'''expanding {name} {amax} -> {mod._weight_quantizer._amax}'''     )







def __UpperCamelCase( _A       :					List[str]     ):
   '''simple docstring'''

   for name, mod in model.named_modules():
      if hasattr(_A	,      '''_weight_quantizer'''     ):
         if not hasattr(mod.weight_quantizer	,      '''_amax'''     ):
            print('''RECALIB: {name:{name_width}} WARNING: NO AMAX BUFFER'''     )
            continue

         # determine which axes to reduce across
         # e.g. a 4D tensor quantized per axis 0 should reduce over (1,2,3)
         UpperCAmelCase__				:					Any 					=	set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis     )
         UpperCAmelCase__				:					Optional[Any] 					=	set(range(len(mod.weight.size()     )     )     ) - axis_set
         UpperCAmelCase__				:					int 					=	pytorch_quantization.utils.reduce_amax(mod.weight	,      axis=_A	,      keepdims=_A     ).detach()
         logger.info(F'''RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}'''     )
         UpperCAmelCase__				:					str 					=	amax







def __UpperCamelCase( _A       :					Dict	,      _A       :					Tuple=25	,      _A       :					Any=1_80	,      _A       :					Optional[int]=None     ):
   '''simple docstring'''

   if ignore is None:
      UpperCAmelCase__				:					Dict 					=	[]
   elif not isinstance(_A	,      _A     ):
      UpperCAmelCase__				:					int 					=	[ignore]

   UpperCAmelCase__				:					Optional[int] 					=	0
   for name, mod in model.named_modules():
      if not hasattr(_A	,      '''weight'''     ):
         continue
      UpperCAmelCase__				:					Dict 					=	max(_A	,      len(_A     )     )

   for name, mod in model.named_modules():
      UpperCAmelCase__				:					str 					=	getattr(_A	,      '''_input_quantizer'''	,      _A     )
      UpperCAmelCase__				:					int 					=	getattr(_A	,      '''_weight_quantizer'''	,      _A     )
      if not hasattr(_A	,      '''weight'''     ):
         continue
      if type(_A     ) in ignore:
         continue
      if [True for s in ignore if type(_A     ) is str and s in name]:
         continue
      UpperCAmelCase__				:					Dict 					=	F'''Act:{input_q.extra_repr()}'''
      UpperCAmelCase__				:					int 					=	F'''Wgt:{weight_q.extra_repr()}'''
      UpperCAmelCase__				:					Dict 					=	F'''{name:{name_width}} {act_str} {wgt_str}'''
      if len(_A     ) <= line_width:
         logger.info(_A     )
      else:
         logger.info(F'''{name:{name_width}} {act_str}'''     )
         logger.info(F'''{'  ':{name_width}} {wgt_str}'''     )







def __UpperCamelCase( _A       :					Optional[Any]     ):
   '''simple docstring'''
   UpperCAmelCase__				:					str 					=	0
   for name, mod in model.named_modules():
      if isinstance(_A	,      pytorch_quantization.nn.TensorQuantizer     ):
         print(F'''{name:80} {mod}'''     )
         count += 1
   print(F'''{count} TensorQuantizers found in model'''     )







def __UpperCamelCase( _A       :					Dict	,      _A       :					Optional[Any]	,      _A       :					Union[str, Any]	,      _A       :					Union[str, Any]	,      _A       :					Dict     ):
   '''simple docstring'''
   UpperCAmelCase__				:					Any 					=	getattr(_A	,      _A	,      _A     )
   if quantizer_mod is not None:
      assert hasattr(_A	,      _A     )
      setattr(_A	,      _A	,      _A     )
   else:
      logger.warning(F'''{name} has no {quantizer}'''     )







def __UpperCamelCase( _A       :					str	,      _A       :					Any	,      _A       :					Optional[int]="both"	,      **_A       :					List[str]     ):
   '''simple docstring'''
   UpperCAmelCase__				:					Tuple 					=	F'''Warning: changing {which} quantizers of {name:{qname_width}}'''
   for k, v in kwargs.items():
      s += F''' {k}={v}'''
      if which in ["input", "both"]:
         set_quantizer(_A	,      _A	,      '''_input_quantizer'''	,      _A	,      _A     )
      if which in ["weight", "both"]:
         set_quantizer(_A	,      _A	,      '''_weight_quantizer'''	,      _A	,      _A     )
   logger.info(_A     )







def __UpperCamelCase( _A       :					Tuple	,      _A       :					List[str]	,      **_A       :					Optional[Any]     ):
   '''simple docstring'''

   for name, mod in model.named_modules():
      if hasattr(_A	,      '''_input_quantizer'''     ) or hasattr(_A	,      '''_weight_quantizer'''     ):
         for n in names:
            if re.search(_A	,      _A     ):
               set_quantizers(_A	,      _A	,      **_A     )
      elif name.endswith('''_quantizer'''     ):
         for n in names:
            if re.search(_A	,      _A     ):
               UpperCAmelCase__				:					str 					=	F'''Warning: changing {name:{name_width}}'''
               for k, v in kwargs.items():
                  s += F''' {k}={v}'''
                  setattr(_A	,      _A	,      _A     )
               logger.info(_A     )




