





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()
_UpperCAmelCase							:			Dict        =							logging.get_logger(__name__)


def 			SCREAMING_SNAKE_CASE (			_UpperCAmelCase					)     ->	Dict:
    lowerCamelCase__       : Any   					=						ASTConfig()

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

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

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

    return config


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

    return name


def 			SCREAMING_SNAKE_CASE (			_UpperCAmelCase		,      _UpperCAmelCase					)     ->	str:
    for key in orig_state_dict.copy().keys():
        lowerCamelCase__       : List[Any]   					=						orig_state_dict.pop(_UpperCAmelCase					)

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

    return orig_state_dict


def 			SCREAMING_SNAKE_CASE (			_UpperCAmelCase					)     ->	List[Any]:
    lowerCamelCase__       : 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(_UpperCAmelCase		,      _UpperCAmelCase					)


@torch.no_grad()
def 			SCREAMING_SNAKE_CASE (			_UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase=False					)     ->	str:
    lowerCamelCase__       : str   					=						get_audio_spectrogram_transformer_config(_UpperCAmelCase					)

    lowerCamelCase__       : List[Any]   					=						{
        '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
    lowerCamelCase__       : Tuple   					=						model_name_to_url[model_name]
    lowerCamelCase__       : Optional[int]   					=						torch.hub.load_state_dict_from_url(_UpperCAmelCase		,      map_location='cpu'					)
    # remove some keys
    remove_keys(_UpperCAmelCase					)
    # rename some keys
    lowerCamelCase__       : List[str]   					=						convert_state_dict(_UpperCAmelCase		,      _UpperCAmelCase					)

    # load 🤗 model
    lowerCamelCase__       : str   					=						ASTForAudioClassification(_UpperCAmelCase					)
    model.eval()

    model.load_state_dict(_UpperCAmelCase					)

    # verify outputs on dummy input
    # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
    lowerCamelCase__       : str   					=						-4.2_677_393 if 'speech-commands' not in model_name else -6.845_978
    lowerCamelCase__       : Union[str, Any]   					=						4.5_689_974 if 'speech-commands' not in model_name else 5.5_654_526
    lowerCamelCase__       : List[str]   					=						1024 if 'speech-commands' not in model_name else 128
    lowerCamelCase__       : List[str]   					=						ASTFeatureExtractor(mean=_UpperCAmelCase		,      std=_UpperCAmelCase		,      max_length=_UpperCAmelCase					)

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

        lowerCamelCase__					,				lowerCamelCase__       : Union[str, Any]   					=						torchaudio.load(_UpperCAmelCase					)
        lowerCamelCase__       : int   					=						waveform.squeeze().numpy()

    lowerCamelCase__       : Optional[int]   					=						feature_extractor(_UpperCAmelCase		,      sampling_rate=1_6000		,      return_tensors='pt'					)

    # forward pass
    lowerCamelCase__       : Union[str, Any]   					=						model(**_UpperCAmelCase					)
    lowerCamelCase__       : Dict   					=						outputs.logits

    if model_name == "ast-finetuned-audioset-10-10-0.4593":
        lowerCamelCase__       : Tuple   					=						torch.tensor([-0.8_760, -7.0_042, -8.6_602]					)
    elif model_name == "ast-finetuned-audioset-10-10-0.450":
        lowerCamelCase__       : Any   					=						torch.tensor([-1.1_986, -7.0_903, -8.2_718]					)
    elif model_name == "ast-finetuned-audioset-10-10-0.448":
        lowerCamelCase__       : Any   					=						torch.tensor([-2.6_128, -8.0_080, -9.4_344]					)
    elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
        lowerCamelCase__       : Tuple   					=						torch.tensor([-1.5_080, -7.4_534, -8.8_917]					)
    elif model_name == "ast-finetuned-audioset-12-12-0.447":
        lowerCamelCase__       : Tuple   					=						torch.tensor([-0.5_050, -6.5_833, -8.0_843]					)
    elif model_name == "ast-finetuned-audioset-14-14-0.443":
        lowerCamelCase__       : List[Any]   					=						torch.tensor([-0.3_826, -7.0_336, -8.2_413]					)
    elif model_name == "ast-finetuned-audioset-16-16-0.442":
        lowerCamelCase__       : Optional[Any]   					=						torch.tensor([-1.2_113, -6.9_101, -8.3_470]					)
    elif model_name == "ast-finetuned-speech-commands-v2":
        lowerCamelCase__       : str   					=						torch.tensor([6.1_589, -8.0_566, -8.7_984]					)
    else:
        raise ValueError('Unknown model name'					)
    if not torch.allclose(logits[0, :3]		,      _UpperCAmelCase		,      atol=1e-4					):
        raise ValueError('Logits don\'t match'					)
    print('Looks ok!'					)

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

    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__":
    _UpperCAmelCase							:			Optional[Any]        =							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."""
    )

    _UpperCAmelCase							:			Any        =							parser.parse_args()
    convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)












# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from ..models.auto import AutoModelForSeqaSeqLM, AutoTokenizer
from .base import PipelineTool

class   lowerCAmelCase				(   __UpperCamelCase     ):
 UpperCAmelCase__						=						"""philschmid/bart-large-cnn-samsum"""
 UpperCAmelCase__						=						(
     """This is a tool that summarizes an English text. It takes an input `text` containing the text to summarize, """
     """and returns a summary of the text."""
 )
 UpperCAmelCase__						=						"""summarizer"""
 UpperCAmelCase__						=						AutoTokenizer
 UpperCAmelCase__						=						AutoModelForSeqaSeqLM

 UpperCAmelCase__						=						["""text"""]
 UpperCAmelCase__						=						["""text"""]
 def 	A_      ( self    : Tuple     ,	UpperCAmelCase    : List[Any]      )     ->      List[str]:
     return self.pre_processor(UpperCAmelCase     ,	return_tensors='pt'     ,	truncation=UpperCAmelCase      )
 def 	A_      ( self    : Optional[int]     ,	UpperCAmelCase    : Dict      )     ->      str:
     return self.model.generate(**UpperCAmelCase      )[0]







 def 	A_      ( self    : Tuple     ,	UpperCAmelCase    : List[Any]      )     ->      Union[str, Any]:
     return self.pre_processor.decode(UpperCAmelCase     ,	skip_special_tokens=UpperCAmelCase     ,	clean_up_tokenization_spaces=UpperCAmelCase      )





