


"""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 __lowerCAmelCase		:





     def __init__(     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase=13     ,					__UpperCAmelCase=2     ,					__UpperCAmelCase=24     ,					__UpperCAmelCase=16     ,					__UpperCAmelCase=True     ,					__UpperCAmelCase=True     ,					__UpperCAmelCase=32     ,					__UpperCAmelCase=5     ,					__UpperCAmelCase=4     ,					__UpperCAmelCase=37     ,					__UpperCAmelCase="gelu"     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=10     ,					__UpperCAmelCase=0.0_2     ,					__UpperCAmelCase=None     ,					__UpperCAmelCase=2     ,					__UpperCAmelCase=2     ,					):

      '''simple docstring'''




      __UpperCamelCase      =	parent
      __UpperCamelCase      =	batch_size
      __UpperCamelCase      =	patch_size
      __UpperCamelCase      =	max_length
      __UpperCamelCase      =	num_mel_bins
      __UpperCamelCase      =	is_training
      __UpperCamelCase      =	use_labels
      __UpperCamelCase      =	hidden_size
      __UpperCamelCase      =	num_hidden_layers
      __UpperCamelCase      =	num_attention_heads
      __UpperCamelCase      =	intermediate_size
      __UpperCamelCase      =	hidden_act
      __UpperCamelCase      =	hidden_dropout_prob
      __UpperCamelCase      =	attention_probs_dropout_prob
      __UpperCamelCase      =	type_sequence_label_size
      __UpperCamelCase      =	initializer_range
      __UpperCamelCase      =	scope
      __UpperCamelCase      =	frequency_stride
      __UpperCamelCase      =	time_stride

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





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase      =	floats_tensor([self.batch_size, self.max_length, self.num_mel_bins]				)

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

      __UpperCamelCase      =	self.get_config()

      return config, input_values, labels





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      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=__UpperCAmelCase     ,					initializer_range=self.initializer_range     ,					frequency_stride=self.frequency_stride     ,					time_stride=self.time_stride     ,					)





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				):

      '''simple docstring'''




      __UpperCamelCase      =	ASTModel(config=__UpperCAmelCase				)
      model.to(__UpperCAmelCase				)
      model.eval()
      __UpperCamelCase      =	model(__UpperCAmelCase				)
      self.parent.assertEqual(result.last_hidden_state.shape     ,					(self.batch_size, self.seq_length, self.hidden_size)				)







     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase      =	self.prepare_config_and_inputs()
      (
          __UpperCamelCase      
      )      =	config_and_inputs
      __UpperCamelCase      =	{'input_values': input_values}
      return config, inputs_dict


@require_torch
class __lowerCAmelCase		(				SCREAMING_SNAKE_CASE__      ,  SCREAMING_SNAKE_CASE__      ,  unittest.TestCase      ):
     lowercase						       =						(
         (
             ASTModel,
             ASTForAudioClassification,
         )
         if is_torch_available()
         else ()
     )
     lowercase						       =						(
         {"audio-classification": ASTForAudioClassification, "feature-extraction": ASTModel}
         if is_torch_available()
         else {}
     )
     lowercase						       =						False
     lowercase						       =						False
     lowercase						       =						False
     lowercase						       =						False





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				):

      '''simple docstring'''




      if pipeline_test_casse_name == "AudioClassificationPipelineTests":
       return True

      return False





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase      =	ASTModelTester(self				)
      __UpperCamelCase      =	ConfigTester(self     ,					config_class=__UpperCAmelCase     ,					has_text_modality=__UpperCAmelCase     ,					hidden_size=37				)





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      self.config_tester.run_common_tests()





     @unittest.skip(reason='AST does not use inputs_embeds'				)
     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      pass





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase      =	self.model_tester.prepare_config_and_inputs_for_common()

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





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase      =	self.model_tester.prepare_config_and_inputs_for_common()

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

       __UpperCamelCase      =	['input_values']
       self.assertListEqual(arg_names[:1]     ,					__UpperCAmelCase				)





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase      =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__UpperCAmelCase				)







     @slow
     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      for model_name in AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       __UpperCamelCase      =	ASTModel.from_pretrained(__UpperCAmelCase				)
       self.assertIsNotNone(__UpperCAmelCase				)






def 							A						(    )				->     Any:
 __UpperCamelCase      =	hf_hub_download(
     repo_id='nielsr/audio-spectogram-transformer-checkpoint'    , filename='sample_audio.flac'    , repo_type='dataset'		)

 __UpperCamelCase      =	torchaudio.load(_UpperCAmelCase		)

 return audio, sampling_rate





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





     @cached_property
     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      return (
          ASTFeatureExtractor.from_pretrained('MIT/ast-finetuned-audioset-10-10-0.4593'				)
          if is_torchaudio_available()
          else None
      )







     @slow
     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase      =	self.default_feature_extractor
      __UpperCamelCase      =	ASTForAudioClassification.from_pretrained('MIT/ast-finetuned-audioset-10-10-0.4593'				).to(__UpperCAmelCase				)

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

      # forward pass
      with torch.no_grad():
       __UpperCamelCase      =	model(**__UpperCAmelCase				)

      # verify the logits
      __UpperCamelCase      =	torch.Size((1, 527)				)
      self.assertEqual(outputs.logits.shape     ,					__UpperCAmelCase				)

      __UpperCamelCase      =	torch.tensor([-0.8_7_6_0, -7.0_0_4_2, -8.6_6_0_2]				).to(__UpperCAmelCase				)

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





"""simple docstring"""

from typing import Union

import fire
import torch
from tqdm import tqdm






def 							A						(    snake_case				:str    , snake_case				:str = "cpu"    , snake_case				:Union[str, None] = None		)				->     None:
 __UpperCamelCase      =	torch.load(snake_case    , map_location=snake_case		)
 for k, v in tqdm(state_dict.items()		):
  if not isinstance(snake_case    , torch.Tensor		):
   raise TypeError('FP16 conversion only works on paths that are saved state dicts, like pytorch_model.bin'		)
  __UpperCamelCase      =	v.half()
 if save_path is None:  # overwrite src_path
  __UpperCamelCase      =	src_path
 torch.save(snake_case    , snake_case		)


if __name__ == "__main__":
      fire.Fire(convert)



