'''simple docstring'''




import json
import os
import unittest

from transformers import BatchEncoding, MvpTokenizer, MvpTokenizerFast
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers, require_torch
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin, filter_roberta_detectors


@require_tokenizers
class       UpperCAmelCase__  (       lowercase_					,     unittest.TestCase			):






   """simple docstring"""
   A							:       Optional[int]     =    MvpTokenizer
   A							:       Optional[int]     =    MvpTokenizerFast
   A							:       Optional[int]     =    True
   A							:       Tuple     =    filter_roberta_detectors
   def 		_lowerCamelCase (self				)    ->     Union[str, Any]:
        super().setUp()
        lowercase_					:       Optional[Any]     				=					[
            'l',
            'o',
            'w',
            'e',
            'r',
            's',
            't',
            'i',
            'd',
            'n',
            '\u0120',
            '\u0120l',
            '\u0120n',
            '\u0120lo',
            '\u0120low',
            'er',
            '\u0120lowest',
            '\u0120newer',
            '\u0120wider',
            '<unk>',
        ]
        lowercase_					:       Optional[int]     				=					dict(zip(_a       ,				range(len(_a				)				)				)				)
        lowercase_					:       List[str]     				=					['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
        lowercase_					:       List[str]     				=					{'unk_token': '<unk>'}

        lowercase_					:       str     				=					os.path.join(self.tmpdirname       ,				VOCAB_FILES_NAMES['vocab_file']				)
        lowercase_					:       List[Any]     				=					os.path.join(self.tmpdirname       ,				VOCAB_FILES_NAMES['merges_file']				)
        with open(self.vocab_file       ,				'w'       ,				encoding='utf-8'				) as fp:
             fp.write(json.dumps(_a				) + '\n'				)
        with open(self.merges_file       ,				'w'       ,				encoding='utf-8'				) as fp:
             fp.write('\n'.join(_a				)				)
   def 		_lowerCamelCase (self       ,				**_a				)    ->     List[str]:
        kwargs.update(self.special_tokens_map				)
        return self.tokenizer_class.from_pretrained(self.tmpdirname       ,				**_a				)
   def 		_lowerCamelCase (self       ,				**_a				)    ->     List[Any]:
        kwargs.update(self.special_tokens_map				)
        return self.rust_tokenizer_class.from_pretrained(self.tmpdirname       ,				**_a				)
   def 		_lowerCamelCase (self       ,				_a				)    ->     Optional[int]:
        return "lower newer", "lower newer"
   @cached_property
   def 		_lowerCamelCase (self				)    ->     Any:
        return MvpTokenizer.from_pretrained('RUCAIBox/mvp'				)
   @cached_property
   def 		_lowerCamelCase (self				)    ->     int:
        return MvpTokenizerFast.from_pretrained('RUCAIBox/mvp'				)
   @require_torch
   def 		_lowerCamelCase (self				)    ->     Tuple:
        lowercase_					:       List[Any]     				=					['A long paragraph for summarization.', 'Another paragraph for summarization.']
        lowercase_					:       int     				=					[0, 250, 251, 17_818, 13, 39_186, 1_938, 4, 2]

        for tokenizer in [self.default_tokenizer, self.default_tokenizer_fast]:
             lowercase_					:       Optional[Any]     				=					tokenizer(_a       ,				max_length=len(_a				)       ,				padding=_a       ,				return_tensors='pt'				)
             self.assertIsInstance(_a       ,				_a				)

             self.assertEqual((2, 9)       ,				batch.input_ids.shape				)
             self.assertEqual((2, 9)       ,				batch.attention_mask.shape				)
             lowercase_					:       Optional[int]     				=					batch.input_ids.tolist()[0]
             self.assertListEqual(_a       ,				_a				)
             # Test that special tokens are reset
   @require_torch
   def 		_lowerCamelCase (self				)    ->     Any:
        lowercase_					:       str     				=					['A long paragraph for summarization.', 'Another paragraph for summarization.']
        for tokenizer in [self.default_tokenizer, self.default_tokenizer_fast]:
             lowercase_					:       Optional[int]     				=					tokenizer(_a       ,				padding=_a       ,				return_tensors='pt'				)
             # check if input_ids are returned and no labels
             self.assertIn('input_ids'       ,				_a				)
             self.assertIn('attention_mask'       ,				_a				)
             self.assertNotIn('labels'       ,				_a				)
             self.assertNotIn('decoder_attention_mask'       ,				_a				)
   @require_torch
   def 		_lowerCamelCase (self				)    ->     Optional[int]:
        lowercase_					:       Optional[int]     				=					[
            'Summary of the text.',
            'Another summary.',
        ]
        for tokenizer in [self.default_tokenizer, self.default_tokenizer_fast]:
             lowercase_					:       str     				=					tokenizer(text_target=_a       ,				max_length=32       ,				padding='max_length'       ,				return_tensors='pt'				)
             self.assertEqual(32       ,				targets['input_ids'].shape[1]				)
   @require_torch
   def 		_lowerCamelCase (self				)    ->     str:
        for tokenizer in [self.default_tokenizer, self.default_tokenizer_fast]:
             lowercase_					:       Any     				=					tokenizer(
                 ['I am a small frog' * 1_024, 'I am a small frog']       ,				padding=_a       ,				truncation=_a       ,				return_tensors='pt'				)
             self.assertIsInstance(_a       ,				_a				)
             self.assertEqual(batch.input_ids.shape       ,				(2, 1_024)				)
   @require_torch
   def 		_lowerCamelCase (self				)    ->     List[str]:
        lowercase_					:       Union[str, Any]     				=					['A long paragraph for summarization.']
        lowercase_					:       Dict     				=					[
            'Summary of the text.',
        ]
        for tokenizer in [self.default_tokenizer, self.default_tokenizer_fast]:
             lowercase_					:       str     				=					tokenizer(_a       ,				text_target=_a       ,				return_tensors='pt'				)
             lowercase_					:       Optional[int]     				=					inputs['input_ids']
             lowercase_					:       Union[str, Any]     				=					inputs['labels']
             self.assertTrue((input_ids[:, 0] == tokenizer.bos_token_id).all().item()				)
             self.assertTrue((labels[:, 0] == tokenizer.bos_token_id).all().item()				)
             self.assertTrue((input_ids[:, -1] == tokenizer.eos_token_id).all().item()				)
             self.assertTrue((labels[:, -1] == tokenizer.eos_token_id).all().item()				)
   def 		_lowerCamelCase (self				)    ->     str:
        pass





   def 		_lowerCamelCase (self				)    ->     Optional[int]:
        for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
             with self.subTest(f'''{tokenizer.__class__.__name__} ({pretrained_name})'''				):
                  lowercase_					:       Tuple     				=					self.rust_tokenizer_class.from_pretrained(_a       ,				**_a				)
                  lowercase_					:       Optional[Any]     				=					self.tokenizer_class.from_pretrained(_a       ,				**_a				)
                  lowercase_					:       Dict     				=					'A, <mask> AllenNLP sentence.'
                  lowercase_					:       List[Any]     				=					tokenizer_r.encode_plus(_a       ,				add_special_tokens=_a       ,				return_token_type_ids=_a				)
                  lowercase_					:       Any     				=					tokenizer_p.encode_plus(_a       ,				add_special_tokens=_a       ,				return_token_type_ids=_a				)

                  # token_type_ids should put 0 everywhere
                  self.assertEqual(sum(tokens_r['token_type_ids']				)       ,				sum(tokens_p['token_type_ids']				)				)

                  # attention_mask should put 1 everywhere, so sum over length should be 1
                  self.assertEqual(
                      sum(tokens_r['attention_mask']				) / len(tokens_r['attention_mask']				)       ,				sum(tokens_p['attention_mask']				) / len(tokens_p['attention_mask']				)       ,				)

                  lowercase_					:       List[Any]     				=					tokenizer_r.convert_ids_to_tokens(tokens_r['input_ids']				)
                  lowercase_					:       int     				=					tokenizer_p.convert_ids_to_tokens(tokens_p['input_ids']				)

                  # Rust correctly handles the space before the mask while python doesnt
                  self.assertSequenceEqual(tokens_p['input_ids']       ,				[0, 250, 6, 50_264, 3_823, 487, 21_992, 3_645, 4, 2]				)
                  self.assertSequenceEqual(tokens_r['input_ids']       ,				[0, 250, 6, 50_264, 3_823, 487, 21_992, 3_645, 4, 2]				)

                  self.assertSequenceEqual(
                      _a       ,				['<s>', 'A', ',', '<mask>', 'ĠAllen', 'N', 'LP', 'Ġsentence', '.', '</s>']				)
                  self.assertSequenceEqual(
                      _a       ,				['<s>', 'A', ',', '<mask>', 'ĠAllen', 'N', 'LP', 'Ġsentence', '.', '</s>']				)





'''simple docstring'''




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()
_A								=   logging.get_logger(__name__)

def _UpperCamelCase   (       SCREAMING_SNAKE_CASE_							):
     lowercase_					:       List[Any]     				=					ASTConfig()

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

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

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

     return config

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

     return name

def _UpperCamelCase   (       SCREAMING_SNAKE_CASE_  ,       SCREAMING_SNAKE_CASE_							):
     for key in orig_state_dict.copy().keys():
          lowercase_					:       List[str]     				=					orig_state_dict.pop(SCREAMING_SNAKE_CASE_							)

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

     return orig_state_dict

def _UpperCamelCase   (       SCREAMING_SNAKE_CASE_							):
     lowercase_					:       List[Any]     				=					[
         '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(SCREAMING_SNAKE_CASE_  ,       SCREAMING_SNAKE_CASE_							)

@torch.no_grad()
def _UpperCamelCase   (       SCREAMING_SNAKE_CASE_  ,       SCREAMING_SNAKE_CASE_  ,       SCREAMING_SNAKE_CASE_=False							):
     lowercase_					:       Dict     				=					get_audio_spectrogram_transformer_config(SCREAMING_SNAKE_CASE_							)

     lowercase_					:       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
     lowercase_					:       Dict     				=					model_name_to_url[model_name]
     lowercase_					:       Optional[Any]     				=					torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_  ,       map_location='cpu'							)
     # remove some keys
     remove_keys(SCREAMING_SNAKE_CASE_							)
     # rename some keys
     lowercase_					:       str     				=					convert_state_dict(SCREAMING_SNAKE_CASE_  ,       SCREAMING_SNAKE_CASE_							)

     # load 🤗 model
     lowercase_					:       Optional[Any]     				=					ASTForAudioClassification(SCREAMING_SNAKE_CASE_							)
     model.eval()

     model.load_state_dict(SCREAMING_SNAKE_CASE_							)

     # verify outputs on dummy input
     # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
     lowercase_					:       Tuple     				=					-4.267_7393 if 'speech-commands' not in model_name else -6.84_5978
     lowercase_					:       str     				=					4.568_9974 if 'speech-commands' not in model_name else 5.565_4526
     lowercase_					:       str     				=					1_024 if 'speech-commands' not in model_name else 128
     lowercase_					:       Dict     				=					ASTFeatureExtractor(mean=SCREAMING_SNAKE_CASE_  ,       std=SCREAMING_SNAKE_CASE_  ,       max_length=SCREAMING_SNAKE_CASE_							)

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

          lowercase_				,lowercase_					:       Union[str, Any]     				=					torchaudio.load(SCREAMING_SNAKE_CASE_							)
          lowercase_					:       str     				=					waveform.squeeze().numpy()

     lowercase_					:       str     				=					feature_extractor(SCREAMING_SNAKE_CASE_  ,       sampling_rate=16_000  ,       return_tensors='pt'							)

     # forward pass
     lowercase_					:       Tuple     				=					model(**SCREAMING_SNAKE_CASE_							)
     lowercase_					:       Tuple     				=					outputs.logits

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

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

     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__":
 _A								=   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.'
 )

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




