




import argparse
import os

import jax as jnp
import numpy as onp
import torch
import torch.nn as nn
from music_spectrogram_diffusion import inference
from tax import checkpoints

from diffusers import DDPMScheduler, OnnxRuntimeModel, SpectrogramDiffusionPipeline
from diffusers.pipelines.spectrogram_diffusion import SpectrogramContEncoder, SpectrogramNotesEncoder, TaFilmDecoder


SCREAMING_SNAKE_CASE__       :      Dict  =   'base_with_context'

def      A      (      _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )	-> int:
       lowerCamelCase   :      List[Any]			 =					nn.Parameter(torch.FloatTensor(weights["token_embedder"]["embedding"]    )    )
       lowerCamelCase   :      List[str]			 =					nn.Parameter(
           torch.FloatTensor(weights["Embed_0"]["embedding"]    ) ,requires_grad=_SCREAMING_SNAKE_CASE    )
       for lyr_num, lyr in enumerate(model.encoders    ):
              lowerCamelCase   :      List[str]			 =					weights[f'''layers_{lyr_num}''']
              lowerCamelCase   :      List[str]			 =					nn.Parameter(
                  torch.FloatTensor(ly_weight["pre_attention_layer_norm"]["scale"]    )    )

              lowerCamelCase   :      Dict			 =					ly_weight["attention"]
              lowerCamelCase   :      Optional[int]			 =					nn.Parameter(torch.FloatTensor(attention_weights["query"]["kernel"].T    )    )
              lowerCamelCase   :      Optional[int]			 =					nn.Parameter(torch.FloatTensor(attention_weights["key"]["kernel"].T    )    )
              lowerCamelCase   :      Tuple			 =					nn.Parameter(torch.FloatTensor(attention_weights["value"]["kernel"].T    )    )
              lowerCamelCase   :      List[Any]			 =					nn.Parameter(torch.FloatTensor(attention_weights["out"]["kernel"].T    )    )

              lowerCamelCase   :      Tuple			 =					nn.Parameter(torch.FloatTensor(ly_weight["pre_mlp_layer_norm"]["scale"]    )    )

              lowerCamelCase   :      Union[str, Any]			 =					nn.Parameter(torch.FloatTensor(ly_weight["mlp"]["wi_0"]["kernel"].T    )    )
              lowerCamelCase   :      Any			 =					nn.Parameter(torch.FloatTensor(ly_weight["mlp"]["wi_1"]["kernel"].T    )    )
              lowerCamelCase   :      Dict			 =					nn.Parameter(torch.FloatTensor(ly_weight["mlp"]["wo"]["kernel"].T    )    )

       lowerCamelCase   :      Any			 =					nn.Parameter(torch.FloatTensor(weights["encoder_norm"]["scale"]    )    )
       return model

def      A      (      _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )	-> List[str]:
       lowerCamelCase   :      Union[str, Any]			 =					nn.Parameter(torch.FloatTensor(weights["input_proj"]["kernel"].T    )    )

       lowerCamelCase   :      Optional[Any]			 =					nn.Parameter(
           torch.FloatTensor(weights["Embed_0"]["embedding"]    ) ,requires_grad=_SCREAMING_SNAKE_CASE    )

       for lyr_num, lyr in enumerate(model.encoders    ):
              lowerCamelCase   :      Optional[Any]			 =					weights[f'''layers_{lyr_num}''']
              lowerCamelCase   :      Union[str, Any]			 =					ly_weight["attention"]

              lowerCamelCase   :      Dict			 =					nn.Parameter(torch.FloatTensor(attention_weights["query"]["kernel"].T    )    )
              lowerCamelCase   :      Optional[int]			 =					nn.Parameter(torch.FloatTensor(attention_weights["key"]["kernel"].T    )    )
              lowerCamelCase   :      Union[str, Any]			 =					nn.Parameter(torch.FloatTensor(attention_weights["value"]["kernel"].T    )    )
              lowerCamelCase   :      str			 =					nn.Parameter(torch.FloatTensor(attention_weights["out"]["kernel"].T    )    )
              lowerCamelCase   :      List[str]			 =					nn.Parameter(
                  torch.FloatTensor(ly_weight["pre_attention_layer_norm"]["scale"]    )    )

              lowerCamelCase   :      List[str]			 =					nn.Parameter(torch.FloatTensor(ly_weight["mlp"]["wi_0"]["kernel"].T    )    )
              lowerCamelCase   :      Tuple			 =					nn.Parameter(torch.FloatTensor(ly_weight["mlp"]["wi_1"]["kernel"].T    )    )
              lowerCamelCase   :      str			 =					nn.Parameter(torch.FloatTensor(ly_weight["mlp"]["wo"]["kernel"].T    )    )
              lowerCamelCase   :      Dict			 =					nn.Parameter(torch.FloatTensor(ly_weight["pre_mlp_layer_norm"]["scale"]    )    )

       lowerCamelCase   :      List[Any]			 =					nn.Parameter(torch.FloatTensor(weights["encoder_norm"]["scale"]    )    )

       return model

def      A      (      _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )	-> Tuple:
       lowerCamelCase   :      List[Any]			 =					nn.Parameter(torch.FloatTensor(weights["time_emb_dense0"]["kernel"].T    )    )
       lowerCamelCase   :      Tuple			 =					nn.Parameter(torch.FloatTensor(weights["time_emb_dense1"]["kernel"].T    )    )

       lowerCamelCase   :      List[Any]			 =					nn.Parameter(
           torch.FloatTensor(weights["Embed_0"]["embedding"]    ) ,requires_grad=_SCREAMING_SNAKE_CASE    )

       lowerCamelCase   :      List[Any]			 =					nn.Parameter(
           torch.FloatTensor(weights["continuous_inputs_projection"]["kernel"].T    )    )

       for lyr_num, lyr in enumerate(model.decoders    ):
              lowerCamelCase   :      int			 =					weights[f'''layers_{lyr_num}''']
              lowerCamelCase   :      Optional[int]			 =					nn.Parameter(
                  torch.FloatTensor(ly_weight["pre_self_attention_layer_norm"]["scale"]    )    )

              lowerCamelCase   :      Any			 =					nn.Parameter(
                  torch.FloatTensor(ly_weight["FiLMLayer_0"]["DenseGeneral_0"]["kernel"].T    )    )

              lowerCamelCase   :      Optional[Any]			 =					ly_weight["self_attention"]
              lowerCamelCase   :      str			 =					nn.Parameter(torch.FloatTensor(attention_weights["query"]["kernel"].T    )    )
              lowerCamelCase   :      int			 =					nn.Parameter(torch.FloatTensor(attention_weights["key"]["kernel"].T    )    )
              lowerCamelCase   :      Dict			 =					nn.Parameter(torch.FloatTensor(attention_weights["value"]["kernel"].T    )    )
              lowerCamelCase   :      List[Any]			 =					nn.Parameter(torch.FloatTensor(attention_weights["out"]["kernel"].T    )    )

              lowerCamelCase   :      str			 =					ly_weight["MultiHeadDotProductAttention_0"]
              lowerCamelCase   :      Dict			 =					nn.Parameter(torch.FloatTensor(attention_weights["query"]["kernel"].T    )    )
              lowerCamelCase   :      int			 =					nn.Parameter(torch.FloatTensor(attention_weights["key"]["kernel"].T    )    )
              lowerCamelCase   :      Union[str, Any]			 =					nn.Parameter(torch.FloatTensor(attention_weights["value"]["kernel"].T    )    )
              lowerCamelCase   :      List[Any]			 =					nn.Parameter(torch.FloatTensor(attention_weights["out"]["kernel"].T    )    )
              lowerCamelCase   :      Any			 =					nn.Parameter(
                  torch.FloatTensor(ly_weight["pre_cross_attention_layer_norm"]["scale"]    )    )

              lowerCamelCase   :      str			 =					nn.Parameter(torch.FloatTensor(ly_weight["pre_mlp_layer_norm"]["scale"]    )    )
              lowerCamelCase   :      str			 =					nn.Parameter(
                  torch.FloatTensor(ly_weight["FiLMLayer_1"]["DenseGeneral_0"]["kernel"].T    )    )
              lowerCamelCase   :      Dict			 =					nn.Parameter(torch.FloatTensor(ly_weight["mlp"]["wi_0"]["kernel"].T    )    )
              lowerCamelCase   :      str			 =					nn.Parameter(torch.FloatTensor(ly_weight["mlp"]["wi_1"]["kernel"].T    )    )
              lowerCamelCase   :      List[Any]			 =					nn.Parameter(torch.FloatTensor(ly_weight["mlp"]["wo"]["kernel"].T    )    )

       lowerCamelCase   :      Any			 =					nn.Parameter(torch.FloatTensor(weights["decoder_norm"]["scale"]    )    )

       lowerCamelCase   :      Optional[Any]			 =					nn.Parameter(torch.FloatTensor(weights["spec_out_dense"]["kernel"].T    )    )

       return model

def      A      (      _SCREAMING_SNAKE_CASE    )	-> int:
       lowerCamelCase   :      str			 =					checkpoints.load_tax_checkpoint(args.checkpoint_path    )
       lowerCamelCase   :      Union[str, Any]			 =					jnp.tree_util.tree_map(onp.array ,_SCREAMING_SNAKE_CASE    )

       lowerCamelCase   :      Union[str, Any]			 =					[
           "from __gin__ import dynamic_registration",
           "from music_spectrogram_diffusion.models.diffusion import diffusion_utils",
           "diffusion_utils.ClassifierFreeGuidanceConfig.eval_condition_weight = 2.0",
           "diffusion_utils.DiffusionConfig.classifier_free_guidance = @diffusion_utils.ClassifierFreeGuidanceConfig()",
       ]

       lowerCamelCase   :      Tuple			 =					os.path.join(args.checkpoint_path ,".." ,"config.gin"    )
       lowerCamelCase   :      Optional[Any]			 =					inference.parse_training_gin_file(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      str			 =					inference.InferenceModel(args.checkpoint_path ,_SCREAMING_SNAKE_CASE    )

       lowerCamelCase   :      Optional[int]			 =					DDPMScheduler(beta_schedule="squaredcos_cap_v2" ,variance_type="fixed_large"    )

       lowerCamelCase   :      Dict			 =					SpectrogramNotesEncoder(
           max_length=synth_model.sequence_length["inputs"] ,vocab_size=synth_model.model.module.config.vocab_size ,d_model=synth_model.model.module.config.emb_dim ,dropout_rate=synth_model.model.module.config.dropout_rate ,num_layers=synth_model.model.module.config.num_encoder_layers ,num_heads=synth_model.model.module.config.num_heads ,d_kv=synth_model.model.module.config.head_dim ,d_ff=synth_model.model.module.config.mlp_dim ,feed_forward_proj="gated-gelu" ,)

       lowerCamelCase   :      List[str]			 =					SpectrogramContEncoder(
           input_dims=synth_model.audio_codec.n_dims ,targets_context_length=synth_model.sequence_length["targets_context"] ,d_model=synth_model.model.module.config.emb_dim ,dropout_rate=synth_model.model.module.config.dropout_rate ,num_layers=synth_model.model.module.config.num_encoder_layers ,num_heads=synth_model.model.module.config.num_heads ,d_kv=synth_model.model.module.config.head_dim ,d_ff=synth_model.model.module.config.mlp_dim ,feed_forward_proj="gated-gelu" ,)

       lowerCamelCase   :      Optional[int]			 =					TaFilmDecoder(
           input_dims=synth_model.audio_codec.n_dims ,targets_length=synth_model.sequence_length["targets_context"] ,max_decoder_noise_time=synth_model.model.module.config.max_decoder_noise_time ,d_model=synth_model.model.module.config.emb_dim ,num_layers=synth_model.model.module.config.num_decoder_layers ,num_heads=synth_model.model.module.config.num_heads ,d_kv=synth_model.model.module.config.head_dim ,d_ff=synth_model.model.module.config.mlp_dim ,dropout_rate=synth_model.model.module.config.dropout_rate ,)

       lowerCamelCase   :      List[str]			 =					load_notes_encoder(ta_checkpoint["target"]["token_encoder"] ,_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      Union[str, Any]			 =					load_continuous_encoder(ta_checkpoint["target"]["continuous_encoder"] ,_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      Optional[Any]			 =					load_decoder(ta_checkpoint["target"]["decoder"] ,_SCREAMING_SNAKE_CASE    )

       lowerCamelCase   :      Any			 =					OnnxRuntimeModel.from_pretrained("kashif/soundstream_mel_decoder"    )

       lowerCamelCase   :      List[Any]			 =					SpectrogramDiffusionPipeline(
           notes_encoder=_SCREAMING_SNAKE_CASE ,continuous_encoder=_SCREAMING_SNAKE_CASE ,decoder=_SCREAMING_SNAKE_CASE ,scheduler=_SCREAMING_SNAKE_CASE ,melgan=_SCREAMING_SNAKE_CASE ,)
       if args.save:
              pipe.save_pretrained(args.output_path    )


if __name__ == "__main__":
       SCREAMING_SNAKE_CASE__       :      List[str]  =   argparse.ArgumentParser()

       parser.add_argument('--output_path', default=None, type=str, required=True, help='Path to the converted model.')
       parser.add_argument(
           '--save', default=True, type=bool, required=False, help='Whether to save the converted model or not.'
       )
       parser.add_argument(
           '--checkpoint_path',
           default=f'''{MODEL}/checkpoint_500000''',
           type=str,
           required=False,
           help='Path to the original jax model checkpoint.',
       )
       SCREAMING_SNAKE_CASE__       :      int  =   parser.parse_args()

       main(args)









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 		_lowercase (     lowercase__   ,	lowercase__=1.0   ,	lowercase__=None   ,	lowercase__=None						):
      if rng is None:
            __lowerCAmelCase					:						Any					=					global_rng

      __lowerCAmelCase					:						str					=					[]
      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   __lowercase      (unittest.TestCase ):
 def __init__(							self      ,							A_      ,							A_=7      ,							A_=400      ,							A_=2000      ,							A_=10      ,							A_=160      ,							A_=8      ,							A_=0.0      ,							A_=4000      ,							A_=False      ,							A_=True      ,							)					->List[str]:



       '''simple docstring'''





       __lowerCAmelCase					:						Optional[int]					=					parent
       __lowerCAmelCase					:						Dict					=					batch_size
       __lowerCAmelCase					:						str					=					min_seq_length
       __lowerCAmelCase					:						int					=					max_seq_length
       __lowerCAmelCase					:						Any					=					(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
       __lowerCAmelCase					:						Any					=					padding_value
       __lowerCAmelCase					:						str					=					sampling_rate
       __lowerCAmelCase					:						Optional[Any]					=					return_attention_mask
       __lowerCAmelCase					:						Optional[Any]					=					do_normalize
       __lowerCAmelCase					:						Optional[Any]					=					feature_size
       __lowerCAmelCase					:						Optional[int]					=					chunk_length
       __lowerCAmelCase					:						Optional[Any]					=					hop_length
 def 	UpperCamelCase__  (							self					)					->Optional[Any]:



       '''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 	UpperCamelCase__  (							self      ,							A_=False      ,							A_=False					)					->Optional[Any]:



       '''simple docstring'''





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

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







@require_torch
@require_torchaudio
class   __lowercase      (_UpperCAmelCase		,   unittest.TestCase ):
 _UpperCamelCase			  =  WhisperFeatureExtractor if is_speech_available() else None
 def 	UpperCamelCase__  (							self					)					->Optional[int]:



       '''simple docstring'''





       __lowerCAmelCase					:						Tuple					=					WhisperFeatureExtractionTester(self					)
 def 	UpperCamelCase__  (							self					)					->int:



       '''simple docstring'''





       __lowerCAmelCase					:						str					=					self.feature_extraction_class(**self.feat_extract_dict					)

       with tempfile.TemporaryDirectory() as tmpdirname:
             __lowerCAmelCase					:						List[str]					=					feat_extract_first.save_pretrained(A_					)[0]
             check_json_file_has_correct_format(A_					)
             __lowerCAmelCase					:						int					=					self.feature_extraction_class.from_pretrained(A_					)

       __lowerCAmelCase					:						Dict					=					feat_extract_first.to_dict()
       __lowerCAmelCase					:						Union[str, Any]					=					feat_extract_second.to_dict()
       __lowerCAmelCase					:						Union[str, Any]					=					feat_extract_first.mel_filters
       __lowerCAmelCase					:						Dict					=					feat_extract_second.mel_filters
       self.assertTrue(np.allclose(A_      ,							A_					)					)
       self.assertEqual(A_      ,							A_					)
 def 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





       __lowerCAmelCase					:						str					=					self.feature_extraction_class(**self.feat_extract_dict					)

       with tempfile.TemporaryDirectory() as tmpdirname:
             __lowerCAmelCase					:						Union[str, Any]					=					os.path.join(A_      ,							'''feat_extract.json'''					)
             feat_extract_first.to_json_file(A_					)
             __lowerCAmelCase					:						List[str]					=					self.feature_extraction_class.from_json_file(A_					)

       __lowerCAmelCase					:						List[str]					=					feat_extract_first.to_dict()
       __lowerCAmelCase					:						Tuple					=					feat_extract_second.to_dict()
       __lowerCAmelCase					:						Any					=					feat_extract_first.mel_filters
       __lowerCAmelCase					:						List[str]					=					feat_extract_second.mel_filters
       self.assertTrue(np.allclose(A_      ,							A_					)					)
       self.assertEqual(A_      ,							A_					)
 def 	UpperCamelCase__  (							self					)					->List[Any]:



       '''simple docstring'''





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

       # Test feature size
       __lowerCAmelCase					:						Tuple					=					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
       __lowerCAmelCase					:						Dict					=					feature_extractor(speech_inputs[0]      ,							return_tensors='''np'''					).input_features
       __lowerCAmelCase					:						List[str]					=					feature_extractor(np_speech_inputs[0]      ,							return_tensors='''np'''					).input_features
       self.assertTrue(np.allclose(A_      ,							A_      ,							atol=1e-3					)					)

       # Test batched
       __lowerCAmelCase					:						Union[str, Any]					=					feature_extractor(A_      ,							return_tensors='''np'''					).input_features
       __lowerCAmelCase					:						List[Any]					=					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.
       __lowerCAmelCase					:						int					=					[floats_list((1, x)					)[0] for x in (800, 800, 800)]
       __lowerCAmelCase					:						Optional[int]					=					np.asarray(A_					)
       __lowerCAmelCase					:						Dict					=					feature_extractor(A_      ,							return_tensors='''np'''					).input_features
       __lowerCAmelCase					:						Any					=					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
       __lowerCAmelCase					:						Optional[int]					=					[floats_list((1, x)					)[0] for x in range(200      ,							(feature_extractor.n_samples + 500)      ,							200					)]
       __lowerCAmelCase					:						Dict					=					[np.asarray(A_					) for speech_input in speech_inputs]

       __lowerCAmelCase					:						Tuple					=					[x[: feature_extractor.n_samples] for x in speech_inputs]
       __lowerCAmelCase					:						Optional[int]					=					[np.asarray(A_					) for speech_input in speech_inputs_truncated]

       __lowerCAmelCase					:						Any					=					feature_extractor(A_      ,							return_tensors='''np'''					).input_features
       __lowerCAmelCase					:						List[str]					=					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 	UpperCamelCase__  (							self					)					->Dict:



       '''simple docstring'''





       import torch

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

       for inputs in [py_speech_inputs, np_speech_inputs]:
             __lowerCAmelCase					:						Tuple					=					feature_extractor.pad([{'''input_features''': inputs}]      ,							return_tensors='''np'''					)
             self.assertTrue(np_processed.input_features.dtype == np.floataa					)
             __lowerCAmelCase					:						int					=					feature_extractor.pad([{'''input_features''': inputs}]      ,							return_tensors='''pt'''					)
             self.assertTrue(pt_processed.input_features.dtype == torch.floataa					)
 def 	UpperCamelCase__  (							self      ,							A_					)					->str:



       '''simple docstring'''





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

       return [x["array"] for x in speech_samples]
 def 	UpperCamelCase__  (							self					)					->int:



       '''simple docstring'''





       __lowerCAmelCase					:						Optional[int]					=					torch.tensor(
           [
               0.1_193, -0.0_946, -0.1_098, -0.0_196, 0.0_225, -0.0_690, -0.1_736, 0.0_951,
               0.0_971, -0.0_817, -0.0_702, 0.0_162, 0.0_260, 0.0_017, -0.0_192, -0.1_678,
               0.0_709, -0.1_867, -0.0_655, -0.0_274, -0.0_234, -0.1_884, -0.0_516, -0.0_554,
               -0.0_274, -0.1_425, -0.1_423, 0.0_837, 0.0_377, -0.0_854
           ]					)
       # fmt: on

       __lowerCAmelCase					:						int					=					self._load_datasamples(1					)
       __lowerCAmelCase					:						Any					=					WhisperFeatureExtractor()
       __lowerCAmelCase					:						Optional[Any]					=					feature_extractor(A_      ,							return_tensors='''pt'''					).input_features
       self.assertEqual(input_features.shape      ,							(1, 80, 3000)					)
       self.assertTrue(torch.allclose(input_features[0, 0, :30]      ,							A_      ,							atol=1e-4					)					)





 def 	UpperCamelCase__  (							self					)					->Any:



       '''simple docstring'''





       __lowerCAmelCase					:						Tuple					=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()					)
       __lowerCAmelCase					:						str					=					self._load_datasamples(1					)[0]
       __lowerCAmelCase					:						Optional[Any]					=					((audio - audio.min()) / (audio.max() - audio.min())) * 6_5535  # Rescale to [0, 65535] to show issue
       __lowerCAmelCase					:						Union[str, Any]					=					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					)					)


