






import argparse
import json
from pathlib import Path

import torch
import torchaudio
from datasets import load_dataset
from huggingface_hub import hf_hub_download

from transformers import ASTConfig, ASTFeatureExtractor, ASTForAudioClassification
from transformers.utils import logging


logging.set_verbosity_info()
SCREAMING_SNAKE_CASE 							=						logging.get_logger(__name__)

def      _lowerCamelCase	(						__A    :							Optional[Any] )      ->					List[Any]:
    _UpperCAmelCase				:      Optional[Any]			  =				ASTConfig()

    if "10-10" in model_name:
        pass
    elif "speech-commands" in model_name:
        _UpperCAmelCase				:      Optional[Any]			  =				128
    elif "12-12" in model_name:
        _UpperCAmelCase				:      Optional[int]			  =				12
        _UpperCAmelCase				:      Optional[int]			  =				12
    elif "14-14" in model_name:
        _UpperCAmelCase				:      Tuple			  =				14
        _UpperCAmelCase				:      Optional[Any]			  =				14
    elif "16-16" in model_name:
        _UpperCAmelCase				:      Union[str, Any]			  =				16
        _UpperCAmelCase				:      List[str]			  =				16
    else:
        raise ValueError('''Model not supported''' )

    _UpperCAmelCase				:      str			  =				'''huggingface/label-files'''
    if "speech-commands" in model_name:
        _UpperCAmelCase				:      Union[str, Any]			  =				35
        _UpperCAmelCase				:      str			  =				'''speech-commands-v2-id2label.json'''
    else:
        _UpperCAmelCase				:      Union[str, Any]			  =				527
        _UpperCAmelCase				:      Optional[int]			  =				'''audioset-id2label.json'''

    _UpperCAmelCase				:      int			  =				json.load(open(hf_hub_download(__A		,		__A		,		repo_type='''dataset''' )		,		'''r''' ) )
    _UpperCAmelCase				:      Union[str, Any]			  =				{int(__A ): v for k, v in idalabel.items()}
    _UpperCAmelCase				:      List[Any]			  =				idalabel
    _UpperCAmelCase				:      Dict			  =				{v: k for k, v in idalabel.items()}

    return config

def      _lowerCamelCase	(						__A    :							Dict )      ->					List[str]:
    if "module.v" in name:
        _UpperCAmelCase				:      Union[str, Any]			  =				name.replace('''module.v'''		,		'''audio_spectrogram_transformer''' )
    if "cls_token" in name:
        _UpperCAmelCase				:      str			  =				name.replace('''cls_token'''		,		'''embeddings.cls_token''' )
    if "dist_token" in name:
        _UpperCAmelCase				:      List[str]			  =				name.replace('''dist_token'''		,		'''embeddings.distillation_token''' )
    if "pos_embed" in name:
        _UpperCAmelCase				:      str			  =				name.replace('''pos_embed'''		,		'''embeddings.position_embeddings''' )
    if "patch_embed.proj" in name:
        _UpperCAmelCase				:      Union[str, Any]			  =				name.replace('''patch_embed.proj'''		,		'''embeddings.patch_embeddings.projection''' )
    # transformer blocks
    if "blocks" in name:
        _UpperCAmelCase				:      int			  =				name.replace('''blocks'''		,		'''encoder.layer''' )
    if "attn.proj" in name:
        _UpperCAmelCase				:      str			  =				name.replace('''attn.proj'''		,		'''attention.output.dense''' )
    if "attn" in name:
        _UpperCAmelCase				:      List[Any]			  =				name.replace('''attn'''		,		'''attention.self''' )
    if "norm1" in name:
        _UpperCAmelCase				:      int			  =				name.replace('''norm1'''		,		'''layernorm_before''' )
    if "norm2" in name:
        _UpperCAmelCase				:      List[str]			  =				name.replace('''norm2'''		,		'''layernorm_after''' )
    if "mlp.fc1" in name:
        _UpperCAmelCase				:      Optional[Any]			  =				name.replace('''mlp.fc1'''		,		'''intermediate.dense''' )
    if "mlp.fc2" in name:
        _UpperCAmelCase				:      Dict			  =				name.replace('''mlp.fc2'''		,		'''output.dense''' )
    # final layernorm
    if "audio_spectrogram_transformer.norm" in name:
        _UpperCAmelCase				:      Optional[Any]			  =				name.replace('''audio_spectrogram_transformer.norm'''		,		'''audio_spectrogram_transformer.layernorm''' )
    # classifier head
    if "module.mlp_head.0" in name:
        _UpperCAmelCase				:      List[Any]			  =				name.replace('''module.mlp_head.0'''		,		'''classifier.layernorm''' )
    if "module.mlp_head.1" in name:
        _UpperCAmelCase				:      List[str]			  =				name.replace('''module.mlp_head.1'''		,		'''classifier.dense''' )

    return name

def      _lowerCamelCase	(						__A    :							int		,		__A    :							Union[str, Any] )      ->					Optional[int]:
    for key in orig_state_dict.copy().keys():
        _UpperCAmelCase				:      List[str]			  =				orig_state_dict.pop(__A )

        if "qkv" in key:
            _UpperCAmelCase				:      int			  =				key.split('''.''' )
            _UpperCAmelCase				:      Optional[Any]			  =				int(key_split[3] )
            _UpperCAmelCase				:      Any			  =				config.hidden_size
            if "weight" in key:
                _UpperCAmelCase				:      str			  =				val[:dim, :]
                _UpperCAmelCase				:      List[Any]			  =				val[dim : dim * 2, :]
                _UpperCAmelCase				:      Union[str, Any]			  =				val[-dim:, :]
            else:
                _UpperCAmelCase				:      List[str]			  =				val[:dim]
                _UpperCAmelCase				:      Tuple			  =				val[dim : dim * 2]
                _UpperCAmelCase				:      List[Any]			  =				val[-dim:]
        else:
            _UpperCAmelCase				:      str			  =				val

    return orig_state_dict

def      _lowerCamelCase	(						__A    :							List[Any] )      ->					Any:
    _UpperCAmelCase				:      Optional[int]			  =				[
        '''module.v.head.weight''',
        '''module.v.head.bias''',
        '''module.v.head_dist.weight''',
        '''module.v.head_dist.bias''',
    ]
    for k in ignore_keys:
        state_dict.pop(__A		,		__A )

@torch.no_grad()
def      _lowerCamelCase	(						__A    :							Union[str, Any]		,		__A    :							Dict		,		__A    :							Any=False )      ->					Tuple:
    _UpperCAmelCase				:      int			  =				get_audio_spectrogram_transformer_config(__A )

    _UpperCAmelCase				:      Optional[int]			  =				{
        '''ast-finetuned-audioset-10-10-0.4593''': (
            '''https://www.dropbox.com/s/ca0b1v2nlxzyeb4/audioset_10_10_0.4593.pth?dl=1'''
        ),
        '''ast-finetuned-audioset-10-10-0.450''': (
            '''https://www.dropbox.com/s/1tv0hovue1bxupk/audioset_10_10_0.4495.pth?dl=1'''
        ),
        '''ast-finetuned-audioset-10-10-0.448''': (
            '''https://www.dropbox.com/s/6u5sikl4b9wo4u5/audioset_10_10_0.4483.pth?dl=1'''
        ),
        '''ast-finetuned-audioset-10-10-0.448-v2''': (
            '''https://www.dropbox.com/s/kt6i0v9fvfm1mbq/audioset_10_10_0.4475.pth?dl=1'''
        ),
        '''ast-finetuned-audioset-12-12-0.447''': (
            '''https://www.dropbox.com/s/snfhx3tizr4nuc8/audioset_12_12_0.4467.pth?dl=1'''
        ),
        '''ast-finetuned-audioset-14-14-0.443''': (
            '''https://www.dropbox.com/s/z18s6pemtnxm4k7/audioset_14_14_0.4431.pth?dl=1'''
        ),
        '''ast-finetuned-audioset-16-16-0.442''': (
            '''https://www.dropbox.com/s/mdsa4t1xmcimia6/audioset_16_16_0.4422.pth?dl=1'''
        ),
        '''ast-finetuned-speech-commands-v2''': (
            '''https://www.dropbox.com/s/q0tbqpwv44pquwy/speechcommands_10_10_0.9812.pth?dl=1'''
        ),
    }

    # load original state_dict
    _UpperCAmelCase				:      Any			  =				model_name_to_url[model_name]
    _UpperCAmelCase				:      Union[str, Any]			  =				torch.hub.load_state_dict_from_url(__A		,		map_location='''cpu''' )
    # remove some keys
    remove_keys(__A )
    # rename some keys
    _UpperCAmelCase				:      Any			  =				convert_state_dict(__A		,		__A )

    # load 🤗 model
    _UpperCAmelCase				:      Optional[Any]			  =				ASTForAudioClassification(__A )
    model.eval()

    model.load_state_dict(__A )

    # verify outputs on dummy input
    # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
    _UpperCAmelCase				:      List[Any]			  =				-4.2677393 if '''speech-commands''' not in model_name else -6.845978
    _UpperCAmelCase				:      Union[str, Any]			  =				4.5689974 if '''speech-commands''' not in model_name else 5.5654526
    _UpperCAmelCase				:      List[Any]			  =				1_024 if '''speech-commands''' not in model_name else 128
    _UpperCAmelCase				:      str			  =				ASTFeatureExtractor(mean=__A		,		std=__A		,		max_length=__A )

    if "speech-commands" in model_name:
        _UpperCAmelCase				:      Any			  =				load_dataset('''speech_commands'''		,		'''v0.02'''		,		split='''validation''' )
        _UpperCAmelCase				:      Dict			  =				dataset[0]['''audio''']['''array''']
    else:
        _UpperCAmelCase				:      Tuple			  =				hf_hub_download(
            repo_id='''nielsr/audio-spectogram-transformer-checkpoint'''		,		filename='''sample_audio.flac'''		,		repo_type='''dataset'''		,		)

        _UpperCAmelCase  , _UpperCAmelCase				:      str			  =				torchaudio.load(__A )
        _UpperCAmelCase				:      Tuple			  =				waveform.squeeze().numpy()

    _UpperCAmelCase				:      str			  =				feature_extractor(__A		,		sampling_rate=16_000		,		return_tensors='''pt''' )

    # forward pass
    _UpperCAmelCase				:      Dict			  =				model(**__A )
    _UpperCAmelCase				:      List[Any]			  =				outputs.logits

    if model_name == "ast-finetuned-audioset-10-10-0.4593":
        _UpperCAmelCase				:      Any			  =				torch.tensor([-0.8760, -7.0042, -8.6602] )
    elif model_name == "ast-finetuned-audioset-10-10-0.450":
        _UpperCAmelCase				:      List[str]			  =				torch.tensor([-1.1986, -7.0903, -8.2718] )
    elif model_name == "ast-finetuned-audioset-10-10-0.448":
        _UpperCAmelCase				:      Optional[int]			  =				torch.tensor([-2.6128, -8.0080, -9.4344] )
    elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
        _UpperCAmelCase				:      List[str]			  =				torch.tensor([-1.5080, -7.4534, -8.8917] )
    elif model_name == "ast-finetuned-audioset-12-12-0.447":
        _UpperCAmelCase				:      List[Any]			  =				torch.tensor([-0.5050, -6.5833, -8.0843] )
    elif model_name == "ast-finetuned-audioset-14-14-0.443":
        _UpperCAmelCase				:      str			  =				torch.tensor([-0.3826, -7.0336, -8.2413] )
    elif model_name == "ast-finetuned-audioset-16-16-0.442":
        _UpperCAmelCase				:      Optional[int]			  =				torch.tensor([-1.2113, -6.9101, -8.3470] )
    elif model_name == "ast-finetuned-speech-commands-v2":
        _UpperCAmelCase				:      Optional[int]			  =				torch.tensor([6.1589, -8.0566, -8.7984] )
    else:
        raise ValueError('''Unknown model name''' )
    if not torch.allclose(logits[0, :3]		,		__A		,		atol=1e-4 ):
        raise ValueError('''Logits don\'t match''' )
    print('''Looks ok!''' )

    if pytorch_dump_folder_path is not None:
        Path(__A ).mkdir(exist_ok=__A )
        print(f'''Saving model {model_name} to {pytorch_dump_folder_path}''' )
        model.save_pretrained(__A )
        print(f'''Saving feature extractor to {pytorch_dump_folder_path}''' )
        feature_extractor.save_pretrained(__A )

    if push_to_hub:
        print('''Pushing model and feature extractor to the hub...''' )
        model.push_to_hub(f'''MIT/{model_name}''' )
        feature_extractor.push_to_hub(f'''MIT/{model_name}''' )


if __name__ == "__main__":
  SCREAMING_SNAKE_CASE 							=						argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--model_name',
      default='ast-finetuned-audioset-10-10-0.4593',
      type=str,
      help='Name of the Audio Spectrogram Transformer model you\'d like to convert.',
  )
  parser.add_argument(
      '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
  )
  parser.add_argument(
      '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
  )

  SCREAMING_SNAKE_CASE 							=						parser.parse_args()
  convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)












from collections import UserDict
from typing import Union

import numpy as np
import requests

from ..utils import (
    add_end_docstrings,
    logging,
)
from .audio_classification import ffmpeg_read
from .base import PIPELINE_INIT_ARGS, Pipeline


SCREAMING_SNAKE_CASE 							=						logging.get_logger(__name__)
@add_end_docstrings(__lowercase			)
class 							A_    (						__lowercase			):






 '''simple docstring'''




 def __init__(	self							,  **_A) ->		int:

     """simple docstring"""



     super().__init__(**_A)

     if self.framework != "pt":
         raise ValueError(f'''The {self.__class__} is only available in PyTorch.''')
     # No specific FOR_XXX available yet


 def __call__(	self							,  _A							,  **_A) ->		List[str]:

     """simple docstring"""



     return super().__call__(_A							,  **_A)


 def  snake_case__				(	self							,  **_A) ->		List[str]:

     """simple docstring"""



     _UpperCAmelCase				:      Optional[int]			  =				{}
     if "candidate_labels" in kwargs:
         _UpperCAmelCase				:      Union[str, Any]			  =				kwargs['''candidate_labels''']
     if "hypothesis_template" in kwargs:
         _UpperCAmelCase				:      Dict			  =				kwargs['''hypothesis_template''']

     return preprocess_params, {}, {}


 def  snake_case__				(	self							,  _A							,  _A=None							,  _A="This is a sound of {}.") ->		Tuple:

     """simple docstring"""



     if isinstance(_A							,  _A):
         if audio.startswith('''http://''') or audio.startswith('''https://'''):
             # We need to actually check for a real protocol, otherwise it's impossible to use a local file
             # like http_huggingface_co.png
             _UpperCAmelCase				:      List[Any]			  =				requests.get(_A).content
         else:
             with open(_A							,  '''rb''') as f:
                 _UpperCAmelCase				:      List[str]			  =				f.read()

     if isinstance(_A							,  _A):
         _UpperCAmelCase				:      List[Any]			  =				ffmpeg_read(_A							,  self.feature_extractor.sampling_rate)

     if not isinstance(_A							,  np.ndarray):
         raise ValueError('''We expect a numpy ndarray as input''')
     if len(audio.shape) != 1:
         raise ValueError('''We expect a single channel audio input for ZeroShotAudioClassificationPipeline''')

     _UpperCAmelCase				:      Optional[Any]			  =				self.feature_extractor(
         [audio]							,  sampling_rate=self.feature_extractor.sampling_rate							,  return_tensors='''pt''')
     _UpperCAmelCase				:      Dict			  =				candidate_labels
     _UpperCAmelCase				:      List[str]			  =				[hypothesis_template.format(_A) for x in candidate_labels]
     _UpperCAmelCase				:      Optional[Any]			  =				self.tokenizer(_A							,  return_tensors=self.framework							,  padding=_A)
     _UpperCAmelCase				:      str			  =				[text_inputs]
     return inputs


 def  snake_case__				(	self							,  _A) ->		List[str]:

     """simple docstring"""



     _UpperCAmelCase				:      Optional[int]			  =				model_inputs.pop('''candidate_labels''')
     _UpperCAmelCase				:      Union[str, Any]			  =				model_inputs.pop('''text_inputs''')
     if isinstance(text_inputs[0]							,  _A):
         _UpperCAmelCase				:      int			  =				text_inputs[0]
     else:
         # Batching case.
         _UpperCAmelCase				:      Dict			  =				text_inputs[0][0]

     _UpperCAmelCase				:      str			  =				self.model(**_A							,  **_A)

     _UpperCAmelCase				:      str			  =				{
         '''candidate_labels''': candidate_labels,
         '''logits''': outputs.logits_per_audio,
     }
     return model_outputs


 def  snake_case__				(	self							,  _A) ->		List[Any]:

     """simple docstring"""



     _UpperCAmelCase				:      str			  =				model_outputs.pop('''candidate_labels''')
     _UpperCAmelCase				:      str			  =				model_outputs['''logits'''][0]

     if self.framework == "pt":
         _UpperCAmelCase				:      int			  =				logits.softmax(dim=0)
         _UpperCAmelCase				:      Optional[Any]			  =				probs.tolist()
     else:
         raise ValueError('''`tf` framework not supported.''')

     _UpperCAmelCase				:      List[str]			  =				[
         {'''score''': score, '''label''': candidate_label}
         for score, candidate_label in sorted(zip(_A							,  _A)							,  key=lambda _A: -x[0])
     ]
     return result




