


'''simple docstring'''






import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ASTConfig
from transformers.testing_utils import require_torch, require_torchaudio, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_torchaudio_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch
		from torch import nn

		from transformers import ASTForAudioClassification, ASTModel
		from transformers.models.audio_spectrogram_transformer.modeling_audio_spectrogram_transformer import (
		    AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		)


if is_torchaudio_available():
		import torchaudio

		from transformers import ASTFeatureExtractor




class 		a__    :





							'''simple docstring'''






							def __init__(		self ,   lowerCamelCase_ ,   lowerCamelCase_=13 ,   lowerCamelCase_=2 ,   lowerCamelCase_=24 ,   lowerCamelCase_=16 ,   lowerCamelCase_=True ,   lowerCamelCase_=True ,   lowerCamelCase_=32 ,   lowerCamelCase_=5 ,   lowerCamelCase_=4 ,   lowerCamelCase_=37 ,   lowerCamelCase_="gelu" ,   lowerCamelCase_=0.1 ,   lowerCamelCase_=0.1 ,   lowerCamelCase_=10 ,   lowerCamelCase_=0.02 ,   lowerCamelCase_=None ,   lowerCamelCase_=2 ,   lowerCamelCase_=2 ,   )	->							Union[str, Any]:
														lowerCAmelCase__	       =     parent
														lowerCAmelCase__	       =     batch_size
														lowerCAmelCase__	       =     patch_size
														lowerCAmelCase__	       =     max_length
														lowerCAmelCase__	       =     num_mel_bins
														lowerCAmelCase__	       =     is_training
														lowerCAmelCase__	       =     use_labels
														lowerCAmelCase__	       =     hidden_size
														lowerCAmelCase__	       =     num_hidden_layers
														lowerCAmelCase__	       =     num_attention_heads
														lowerCAmelCase__	       =     intermediate_size
														lowerCAmelCase__	       =     hidden_act
														lowerCAmelCase__	       =     hidden_dropout_prob
														lowerCAmelCase__	       =     attention_probs_dropout_prob
														lowerCAmelCase__	       =     type_sequence_label_size
														lowerCAmelCase__	       =     initializer_range
														lowerCAmelCase__	       =     scope
														lowerCAmelCase__	       =     frequency_stride
														lowerCAmelCase__	       =     time_stride

														# in AST, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
														lowerCAmelCase__	       =     (self.num_mel_bins - self.patch_size) // self.frequency_stride + 1
														lowerCAmelCase__	       =     (self.max_length - self.patch_size) // self.time_stride + 1
														lowerCAmelCase__	       =     frequency_out_dimension * time_out_dimension
														lowerCAmelCase__	       =     num_patches + 2







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Union[str, Any]:
														lowerCAmelCase__	       =     floats_tensor([self.batch_size, self.max_length, self.num_mel_bins]       )

														lowerCAmelCase__	       =     None
														if self.use_labels:
																					lowerCAmelCase__	       =     ids_tensor([self.batch_size] ,   self.type_sequence_label_size       )

														lowerCAmelCase__	       =     self.get_config()

														return config, input_values, labels







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[int]:
														return ASTConfig(
														    patch_size=self.patch_size ,   max_length=self.max_length ,   num_mel_bins=self.num_mel_bins ,   hidden_size=self.hidden_size ,   num_hidden_layers=self.num_hidden_layers ,   num_attention_heads=self.num_attention_heads ,   intermediate_size=self.intermediate_size ,   hidden_act=self.hidden_act ,   hidden_dropout_prob=self.hidden_dropout_prob ,   attention_probs_dropout_prob=self.attention_probs_dropout_prob ,   is_decoder=lowerCamelCase_ ,   initializer_range=self.initializer_range ,   frequency_stride=self.frequency_stride ,   time_stride=self.time_stride ,   )







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							Optional[Any]:
														lowerCAmelCase__	       =     ASTModel(config=lowerCamelCase_       )
														model.to(lowerCamelCase_       )
														model.eval()
														lowerCAmelCase__	       =     model(lowerCamelCase_       )
														self.parent.assertEqual(result.last_hidden_state.shape ,   (self.batch_size, self.seq_length, self.hidden_size)       )







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							List[Any]:
														lowerCAmelCase__	       =     self.prepare_config_and_inputs()
														(
														    (
														    lowerCAmelCase__
														)       , (
														    lowerCAmelCase__
														)       , (
														    lowerCAmelCase__
														)       , 
														)	       =     config_and_inputs
														lowerCAmelCase__	       =     {'''input_values''': input_values}
														return config, inputs_dict




@require_torch
class 		a__    (							a__    ,    a__    ,    unittest.TestCase  ):





							'''simple docstring'''






							lowercase__      :       Union[str, Any]     						=				(
							    (
							        ASTModel,
							        ASTForAudioClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							lowercase__      :       int     						=				(
							    {"audio-classification": ASTForAudioClassification, "feature-extraction": ASTModel}
							    if is_torch_available()
							    else {}
							)
							lowercase__      :       Optional[Any]     						=				False
							lowercase__      :       Optional[Any]     						=				False
							lowercase__      :       List[Any]     						=				False
							lowercase__      :       str     						=				False







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_       )	->							List[Any]:
														if pipeline_test_casse_name == "AudioClassificationPipelineTests":
																					return True

														return False







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Union[str, Any]:
														lowerCAmelCase__	       =     ASTModelTester(self       )
														lowerCAmelCase__	       =     ConfigTester(self ,   config_class=lowerCamelCase_ ,   has_text_modality=lowerCamelCase_ ,   hidden_size=37       )







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Union[str, Any]:
														self.config_tester.run_common_tests()







							@unittest.skip(reason='''AST does not use inputs_embeds'''       )
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Union[str, Any]:
														pass







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[Any]:
														lowerCAmelCase__       , lowerCAmelCase__	       =     self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					lowerCAmelCase__	       =     model_class(lowerCamelCase_       )
																					self.assertIsInstance(model.get_input_embeddings() ,   (nn.Module)       )
																					lowerCAmelCase__	       =     model.get_output_embeddings()
																					self.assertTrue(x is None or isinstance(lowerCamelCase_ ,   nn.Linear       )       )







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							List[Any]:
														lowerCAmelCase__       , lowerCAmelCase__	       =     self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					lowerCAmelCase__	       =     model_class(lowerCamelCase_       )
																					lowerCAmelCase__	       =     inspect.signature(model.forward       )
																					# signature.parameters is an OrderedDict => so arg_names order is deterministic
																					lowerCAmelCase__	       =     [*signature.parameters.keys()]

																					lowerCAmelCase__	       =     ['''input_values''']
																					self.assertListEqual(arg_names[:1] ,   lowerCamelCase_       )







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Union[str, Any]:
														lowerCAmelCase__	       =     self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(*lowerCamelCase_       )







							@slow
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							List[Any]:
														for model_name in AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					lowerCAmelCase__	       =     ASTModel.from_pretrained(lowerCamelCase_       )
																					self.assertIsNotNone(lowerCamelCase_       )


def  _snake_case       (					)						->   List[Any]:
							lowerCAmelCase__	       =     hf_hub_download(
							    repo_id='''nielsr/audio-spectogram-transformer-checkpoint'''			,   filename='''sample_audio.flac'''			,   repo_type='''dataset'''	)

							lowerCAmelCase__       , lowerCAmelCase__	       =     torchaudio.load(A	)

							return audio, sampling_rate






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





							'''simple docstring'''






							@cached_property
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Tuple:
														return (
														    ASTFeatureExtractor.from_pretrained('''MIT/ast-finetuned-audioset-10-10-0.4593'''       )
														    if is_torchaudio_available()
														    else None
														)







							@slow
							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Dict:
														lowerCAmelCase__	       =     self.default_feature_extractor
														lowerCAmelCase__	       =     ASTForAudioClassification.from_pretrained('''MIT/ast-finetuned-audioset-10-10-0.4593'''       ).to(lowerCamelCase_       )

														lowerCAmelCase__	       =     self.default_feature_extractor
														lowerCAmelCase__       , lowerCAmelCase__	       =     prepare_audio()
														lowerCAmelCase__	       =     audio.squeeze().numpy()
														lowerCAmelCase__	       =     feature_extractor(lowerCamelCase_ ,   sampling_rate=lowerCamelCase_ ,   return_tensors='''pt'''       ).to(lowerCamelCase_       )

														# forward pass
														with torch.no_grad():
																					lowerCAmelCase__	       =     model(**lowerCamelCase_       )

														# verify the logits
														lowerCAmelCase__	       =     torch.Size((1, 5_27)       )
														self.assertEqual(outputs.logits.shape ,   lowerCamelCase_       )

														lowerCAmelCase__	       =     torch.tensor([-0.8_760, -7.0_042, -8.6_602]       ).to(lowerCamelCase_       )

														self.assertTrue(torch.allclose(outputs.logits[0, :3] ,   lowerCamelCase_ ,   atol=1e-4       )       )



'''simple docstring'''






from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import get_down_block, get_mid_block, get_out_block, get_up_block




@dataclass
class 		a__    (							a__  ):





							'''simple docstring'''






							lowercase__      :       torch.FloatTensor




class 		a__    (							a__    ,    a__  ):





							'''simple docstring'''






							@register_to_config
							def __init__(		self ,   lowerCamelCase_ = 6_55_36 ,   lowerCamelCase_ = None ,   lowerCamelCase_ = 2 ,   lowerCamelCase_ = 2 ,   lowerCamelCase_ = 0 ,   lowerCamelCase_ = "fourier" ,   lowerCamelCase_ = True ,   lowerCamelCase_ = False ,   lowerCamelCase_ = 0.0 ,   lowerCamelCase_ = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D") ,   lowerCamelCase_ = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip") ,   lowerCamelCase_ = "UNetMidBlock1D" ,   lowerCamelCase_ = None ,   lowerCamelCase_ = (32, 32, 64) ,   lowerCamelCase_ = None ,   lowerCamelCase_ = 8 ,   lowerCamelCase_ = 1 ,   lowerCamelCase_ = False ,   )	->							Optional[int]:
														super().__init__()
														lowerCAmelCase__	       =     sample_size

														# time
														if time_embedding_type == "fourier":
																					lowerCAmelCase__	       =     GaussianFourierProjection(
																					    embedding_size=8 ,   set_W_to_weight=lowerCamelCase_ ,   log=lowerCamelCase_ ,   flip_sin_to_cos=lowerCamelCase_       )
																					lowerCAmelCase__	       =     2 * block_out_channels[0]
														elif time_embedding_type == "positional":
																					lowerCAmelCase__	       =     Timesteps(
																					    block_out_channels[0] ,   flip_sin_to_cos=lowerCamelCase_ ,   downscale_freq_shift=lowerCamelCase_       )
																					lowerCAmelCase__	       =     block_out_channels[0]

														if use_timestep_embedding:
																					lowerCAmelCase__	       =     block_out_channels[0] * 4
																					lowerCAmelCase__	       =     TimestepEmbedding(
																					    in_channels=lowerCamelCase_ ,   time_embed_dim=lowerCamelCase_ ,   act_fn=lowerCamelCase_ ,   out_dim=block_out_channels[0] ,   )

														lowerCAmelCase__	       =     nn.ModuleList([]       )
														lowerCAmelCase__	       =     None
														lowerCAmelCase__	       =     nn.ModuleList([]       )
														lowerCAmelCase__	       =     None

														# down
														lowerCAmelCase__	       =     in_channels
														for i, down_block_type in enumerate(lowerCamelCase_       ):
																					lowerCAmelCase__	       =     output_channel
																					lowerCAmelCase__	       =     block_out_channels[i]

																					if i == 0:
																												input_channel += extra_in_channels

																					lowerCAmelCase__	       =     i == len(lowerCamelCase_       ) - 1

																					lowerCAmelCase__	       =     get_down_block(
																					    lowerCamelCase_ ,   num_layers=lowerCamelCase_ ,   in_channels=lowerCamelCase_ ,   out_channels=lowerCamelCase_ ,   temb_channels=block_out_channels[0] ,   add_downsample=not is_final_block or downsample_each_block ,   )
																					self.down_blocks.append(lowerCamelCase_       )

														# mid
														lowerCAmelCase__	       =     get_mid_block(
														    lowerCamelCase_ ,   in_channels=block_out_channels[-1] ,   mid_channels=block_out_channels[-1] ,   out_channels=block_out_channels[-1] ,   embed_dim=block_out_channels[0] ,   num_layers=lowerCamelCase_ ,   add_downsample=lowerCamelCase_ ,   )

														# up
														lowerCAmelCase__	       =     list(reversed(lowerCamelCase_       )       )
														lowerCAmelCase__	       =     reversed_block_out_channels[0]
														if out_block_type is None:
																					lowerCAmelCase__	       =     out_channels
														else:
																					lowerCAmelCase__	       =     block_out_channels[0]

														for i, up_block_type in enumerate(lowerCamelCase_       ):
																					lowerCAmelCase__	       =     output_channel
																					lowerCAmelCase__	       =     (
																					    reversed_block_out_channels[i + 1] if i < len(lowerCamelCase_       ) - 1 else final_upsample_channels
																					)

																					lowerCAmelCase__	       =     i == len(lowerCamelCase_       ) - 1

																					lowerCAmelCase__	       =     get_up_block(
																					    lowerCamelCase_ ,   num_layers=lowerCamelCase_ ,   in_channels=lowerCamelCase_ ,   out_channels=lowerCamelCase_ ,   temb_channels=block_out_channels[0] ,   add_upsample=not is_final_block ,   )
																					self.up_blocks.append(lowerCamelCase_       )
																					lowerCAmelCase__	       =     output_channel

														# out
														lowerCAmelCase__	       =     norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4 ,   32       )
														lowerCAmelCase__	       =     get_out_block(
														    out_block_type=lowerCamelCase_ ,   num_groups_out=lowerCamelCase_ ,   embed_dim=block_out_channels[0] ,   out_channels=lowerCamelCase_ ,   act_fn=lowerCamelCase_ ,   fc_dim=block_out_channels[-1] // 4 ,   )







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_ = True ,   )	->							Union[UNetaDOutput, Tuple]:
														lowerCAmelCase__	       =     timestep
														if not torch.is_tensor(lowerCamelCase_       ):
																					lowerCAmelCase__	       =     torch.tensor([timesteps] ,   dtype=torch.long ,   device=sample.device       )
														elif torch.is_tensor(lowerCamelCase_       ) and len(timesteps.shape       ) == 0:
																					lowerCAmelCase__	       =     timesteps[None].to(sample.device       )

														lowerCAmelCase__	       =     self.time_proj(lowerCamelCase_       )
														if self.config.use_timestep_embedding:
																					lowerCAmelCase__	       =     self.time_mlp(lowerCamelCase_       )
														else:
																					lowerCAmelCase__	       =     timestep_embed[..., None]
																					lowerCAmelCase__	       =     timestep_embed.repeat([1, 1, sample.shape[2]]       ).to(sample.dtype       )
																					lowerCAmelCase__	       =     timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:])       )

														# 2. down
														lowerCAmelCase__	       =     ()
														for downsample_block in self.down_blocks:
																					lowerCAmelCase__       , lowerCAmelCase__	       =     downsample_block(hidden_states=lowerCamelCase_ ,   temb=lowerCamelCase_       )
																					down_block_res_samples += res_samples

														# 3. mid
														if self.mid_block:
																					lowerCAmelCase__	       =     self.mid_block(lowerCamelCase_ ,   lowerCamelCase_       )

														# 4. up
														for i, upsample_block in enumerate(self.up_blocks       ):
																					lowerCAmelCase__	       =     down_block_res_samples[-1:]
																					lowerCAmelCase__	       =     down_block_res_samples[:-1]
																					lowerCAmelCase__	       =     upsample_block(lowerCamelCase_ ,   res_hidden_states_tuple=lowerCamelCase_ ,   temb=lowerCamelCase_       )

														# 5. post-process
														if self.out_block:
																					lowerCAmelCase__	       =     self.out_block(lowerCamelCase_ ,   lowerCamelCase_       )

														if not return_dict:
																					return (sample,)

														return UNetaDOutput(sample=lowerCamelCase_       )