

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





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					ASTConfig()

     if "10-10" in model_name:
          pass
     elif "speech-commands" in model_name:
          SCREAMING_SNAKE_CASE__					=					1_28
     elif "12-12" in model_name:
          SCREAMING_SNAKE_CASE__					=					12
          SCREAMING_SNAKE_CASE__					=					12
     elif "14-14" in model_name:
          SCREAMING_SNAKE_CASE__					=					14
          SCREAMING_SNAKE_CASE__					=					14
     elif "16-16" in model_name:
          SCREAMING_SNAKE_CASE__					=					16
          SCREAMING_SNAKE_CASE__					=					16
     else:
          raise ValueError('''Model not supported'''       )

     SCREAMING_SNAKE_CASE__					=					'''huggingface/label-files'''
     if "speech-commands" in model_name:
          SCREAMING_SNAKE_CASE__					=					35
          SCREAMING_SNAKE_CASE__					=					'''speech-commands-v2-id2label.json'''
     else:
          SCREAMING_SNAKE_CASE__					=					5_27
          SCREAMING_SNAKE_CASE__					=					'''audioset-id2label.json'''

     SCREAMING_SNAKE_CASE__					=					json.load(open(hf_hub_download(_A   ,	_A   ,	repo_type='''dataset'''       )   ,	'''r'''       )       )
     SCREAMING_SNAKE_CASE__					=					{int(_A       ): v for k, v in idalabel.items()}
     SCREAMING_SNAKE_CASE__					=					idalabel
     SCREAMING_SNAKE_CASE__					=					{v: k for k, v in idalabel.items()}

     return config
def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     if "module.v" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''module.v'''   ,	'''audio_spectrogram_transformer'''       )
     if "cls_token" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''cls_token'''   ,	'''embeddings.cls_token'''       )
     if "dist_token" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''dist_token'''   ,	'''embeddings.distillation_token'''       )
     if "pos_embed" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''pos_embed'''   ,	'''embeddings.position_embeddings'''       )
     if "patch_embed.proj" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''patch_embed.proj'''   ,	'''embeddings.patch_embeddings.projection'''       )
     # transformer blocks
     if "blocks" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''blocks'''   ,	'''encoder.layer'''       )
     if "attn.proj" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''attn.proj'''   ,	'''attention.output.dense'''       )
     if "attn" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''attn'''   ,	'''attention.self'''       )
     if "norm1" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''norm1'''   ,	'''layernorm_before'''       )
     if "norm2" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''norm2'''   ,	'''layernorm_after'''       )
     if "mlp.fc1" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''mlp.fc1'''   ,	'''intermediate.dense'''       )
     if "mlp.fc2" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''mlp.fc2'''   ,	'''output.dense'''       )
     # final layernorm
     if "audio_spectrogram_transformer.norm" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''audio_spectrogram_transformer.norm'''   ,	'''audio_spectrogram_transformer.layernorm'''       )
     # classifier head
     if "module.mlp_head.0" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''module.mlp_head.0'''   ,	'''classifier.layernorm'''       )
     if "module.mlp_head.1" in name:
          SCREAMING_SNAKE_CASE__					=					name.replace('''module.mlp_head.1'''   ,	'''classifier.dense'''       )

     return name
def        UpperCAmelCase_     (		_A   ,	_A       ):





     '''simple docstring'''




     for key in orig_state_dict.copy().keys():
          SCREAMING_SNAKE_CASE__					=					orig_state_dict.pop(_A       )

          if "qkv" in key:
               SCREAMING_SNAKE_CASE__					=					key.split('''.'''       )
               SCREAMING_SNAKE_CASE__					=					int(key_split[3]       )
               SCREAMING_SNAKE_CASE__					=					config.hidden_size
               if "weight" in key:
                    SCREAMING_SNAKE_CASE__					=					val[:dim, :]
                    SCREAMING_SNAKE_CASE__					=					val[dim : dim * 2, :]
                    SCREAMING_SNAKE_CASE__					=					val[-dim:, :]
               else:
                    SCREAMING_SNAKE_CASE__					=					val[:dim]
                    SCREAMING_SNAKE_CASE__					=					val[dim : dim * 2]
                    SCREAMING_SNAKE_CASE__					=					val[-dim:]
          else:
               SCREAMING_SNAKE_CASE__					=					val

     return orig_state_dict
def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					[
         '''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        UpperCAmelCase_     (		_A   ,	_A   ,	_A=False       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					get_audio_spectrogram_transformer_config(_A       )

     SCREAMING_SNAKE_CASE__					=					{
         '''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
     SCREAMING_SNAKE_CASE__					=					model_name_to_url[model_name]
     SCREAMING_SNAKE_CASE__					=					torch.hub.load_state_dict_from_url(_A   ,	map_location='''cpu'''       )
     # remove some keys
     remove_keys(_A       )
     # rename some keys
     SCREAMING_SNAKE_CASE__					=					convert_state_dict(_A   ,	_A       )

     # load 🤗 model
     SCREAMING_SNAKE_CASE__					=					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
     SCREAMING_SNAKE_CASE__					=					-4.2_6_7_7_3_9_3 if '''speech-commands''' not in model_name else -6.8_4_5_9_7_8
     SCREAMING_SNAKE_CASE__					=					4.5_6_8_9_9_7_4 if '''speech-commands''' not in model_name else 5.5_6_5_4_5_2_6
     SCREAMING_SNAKE_CASE__					=					10_24 if '''speech-commands''' not in model_name else 1_28
     SCREAMING_SNAKE_CASE__					=					ASTFeatureExtractor(mean=_A   ,	std=_A   ,	max_length=_A       )

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

          SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__					=					torchaudio.load(_A       )
          SCREAMING_SNAKE_CASE__					=					waveform.squeeze().numpy()

     SCREAMING_SNAKE_CASE__					=					feature_extractor(_A   ,	sampling_rate=1_60_00   ,	return_tensors='''pt'''       )

     # forward pass
     SCREAMING_SNAKE_CASE__					=					model(**_A       )
     SCREAMING_SNAKE_CASE__					=					outputs.logits

     if model_name == "ast-finetuned-audioset-10-10-0.4593":
          SCREAMING_SNAKE_CASE__					=					torch.tensor([-0.8_7_6_0, -7.0_0_4_2, -8.6_6_0_2]       )
     elif model_name == "ast-finetuned-audioset-10-10-0.450":
          SCREAMING_SNAKE_CASE__					=					torch.tensor([-1.1_9_8_6, -7.0_9_0_3, -8.2_7_1_8]       )
     elif model_name == "ast-finetuned-audioset-10-10-0.448":
          SCREAMING_SNAKE_CASE__					=					torch.tensor([-2.6_1_2_8, -8.0_0_8_0, -9.4_3_4_4]       )
     elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
          SCREAMING_SNAKE_CASE__					=					torch.tensor([-1.5_0_8_0, -7.4_5_3_4, -8.8_9_1_7]       )
     elif model_name == "ast-finetuned-audioset-12-12-0.447":
          SCREAMING_SNAKE_CASE__					=					torch.tensor([-0.5_0_5_0, -6.5_8_3_3, -8.0_8_4_3]       )
     elif model_name == "ast-finetuned-audioset-14-14-0.443":
          SCREAMING_SNAKE_CASE__					=					torch.tensor([-0.3_8_2_6, -7.0_3_3_6, -8.2_4_1_3]       )
     elif model_name == "ast-finetuned-audioset-16-16-0.442":
          SCREAMING_SNAKE_CASE__					=					torch.tensor([-1.2_1_1_3, -6.9_1_0_1, -8.3_4_7_0]       )
     elif model_name == "ast-finetuned-speech-commands-v2":
          SCREAMING_SNAKE_CASE__					=					torch.tensor([6.1_5_8_9, -8.0_5_6_6, -8.7_9_8_4]       )
     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  :						Optional[int]         =    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  :						str         =    parser.parse_args()
       convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





import argparse
import json
from collections import OrderedDict
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import PoolFormerConfig, PoolFormerForImageClassification, PoolFormerImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
_SCREAMING_SNAKE_CASE  :						Dict         =    logging.get_logger(__name__)
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A   ,	_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					original_name.split('''.'''       )[0]
     SCREAMING_SNAKE_CASE__					=					key.split('''.'''       )
     SCREAMING_SNAKE_CASE__					=					int(key_list[key_list.index(_A       ) - 2]       )
     SCREAMING_SNAKE_CASE__					=					int(key_list[key_list.index(_A       ) - 1]       )
     SCREAMING_SNAKE_CASE__					=					orig_block_num - offset

     SCREAMING_SNAKE_CASE__					=					key.replace(F'''{orig_block_num}.{layer_num}.{original_name}'''   ,	F'''block.{new_block_num}.{layer_num}.{new_name}'''       )
     return key
def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					OrderedDict()
     SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__					=					0, 0
     for key, value in state_dict.items():
          if key.startswith('''network'''       ):
               SCREAMING_SNAKE_CASE__					=					key.replace('''network'''   ,	'''poolformer.encoder'''       )
          if "proj" in key:
               # Works for the first embedding as well as the internal embedding layers
               if key.endswith('''bias'''       ) and "patch_embed" not in key:
                    patch_emb_offset += 1
               SCREAMING_SNAKE_CASE__					=					key[: key.find('''proj'''       )]
               SCREAMING_SNAKE_CASE__					=					key.replace(_A   ,	F'''patch_embeddings.{total_embed_found}.'''       )
               SCREAMING_SNAKE_CASE__					=					key.replace('''proj'''   ,	'''projection'''       )
               if key.endswith('''bias'''       ):
                    total_embed_found += 1
          if "patch_embeddings" in key:
               SCREAMING_SNAKE_CASE__					=					'''poolformer.encoder.''' + key
          if "mlp.fc1" in key:
               SCREAMING_SNAKE_CASE__					=					replace_key_with_offset(_A   ,	_A   ,	'''mlp.fc1'''   ,	'''output.conv1'''       )
          if "mlp.fc2" in key:
               SCREAMING_SNAKE_CASE__					=					replace_key_with_offset(_A   ,	_A   ,	'''mlp.fc2'''   ,	'''output.conv2'''       )
          if "norm1" in key:
               SCREAMING_SNAKE_CASE__					=					replace_key_with_offset(_A   ,	_A   ,	'''norm1'''   ,	'''before_norm'''       )
          if "norm2" in key:
               SCREAMING_SNAKE_CASE__					=					replace_key_with_offset(_A   ,	_A   ,	'''norm2'''   ,	'''after_norm'''       )
          if "layer_scale_1" in key:
               SCREAMING_SNAKE_CASE__					=					replace_key_with_offset(_A   ,	_A   ,	'''layer_scale_1'''   ,	'''layer_scale_1'''       )
          if "layer_scale_2" in key:
               SCREAMING_SNAKE_CASE__					=					replace_key_with_offset(_A   ,	_A   ,	'''layer_scale_2'''   ,	'''layer_scale_2'''       )
          if "head" in key:
               SCREAMING_SNAKE_CASE__					=					key.replace('''head'''   ,	'''classifier'''       )
          SCREAMING_SNAKE_CASE__					=					value
     return new_state_dict
def        UpperCAmelCase_     (		):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					'''http://images.cocodataset.org/val2017/000000039769.jpg'''
     SCREAMING_SNAKE_CASE__					=					Image.open(requests.get(_A   ,	stream=_A       ).raw       )

     return image
@torch.no_grad()
def        UpperCAmelCase_     (		_A   ,	_A   ,	_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					PoolFormerConfig()

     # set attributes based on model_name
     SCREAMING_SNAKE_CASE__					=					'''huggingface/label-files'''
     SCREAMING_SNAKE_CASE__					=					model_name[-3:]
     SCREAMING_SNAKE_CASE__					=					10_00
     SCREAMING_SNAKE_CASE__					=					'''imagenet-1k-id2label.json'''
     SCREAMING_SNAKE_CASE__					=					(1, 10_00)

     # set config attributes
     SCREAMING_SNAKE_CASE__					=					json.load(open(hf_hub_download(_A   ,	_A   ,	repo_type='''dataset'''       )   ,	'''r'''       )       )
     SCREAMING_SNAKE_CASE__					=					{int(_A       ): v for k, v in idalabel.items()}
     SCREAMING_SNAKE_CASE__					=					idalabel
     SCREAMING_SNAKE_CASE__					=					{v: k for k, v in idalabel.items()}
     if size == "s12":
          SCREAMING_SNAKE_CASE__					=					[2, 2, 6, 2]
          SCREAMING_SNAKE_CASE__					=					[64, 1_28, 3_20, 5_12]
          SCREAMING_SNAKE_CASE__					=					4.0
          SCREAMING_SNAKE_CASE__					=					0.9
     elif size == "s24":
          SCREAMING_SNAKE_CASE__					=					[4, 4, 12, 4]
          SCREAMING_SNAKE_CASE__					=					[64, 1_28, 3_20, 5_12]
          SCREAMING_SNAKE_CASE__					=					4.0
          SCREAMING_SNAKE_CASE__					=					0.9
     elif size == "s36":
          SCREAMING_SNAKE_CASE__					=					[6, 6, 18, 6]
          SCREAMING_SNAKE_CASE__					=					[64, 1_28, 3_20, 5_12]
          SCREAMING_SNAKE_CASE__					=					4.0
          SCREAMING_SNAKE_CASE__					=					1e-6
          SCREAMING_SNAKE_CASE__					=					0.9
     elif size == "m36":
          SCREAMING_SNAKE_CASE__					=					[6, 6, 18, 6]
          SCREAMING_SNAKE_CASE__					=					[96, 1_92, 3_84, 7_68]
          SCREAMING_SNAKE_CASE__					=					4.0
          SCREAMING_SNAKE_CASE__					=					1e-6
          SCREAMING_SNAKE_CASE__					=					0.9_5
     elif size == "m48":
          SCREAMING_SNAKE_CASE__					=					[8, 8, 24, 8]
          SCREAMING_SNAKE_CASE__					=					[96, 1_92, 3_84, 7_68]
          SCREAMING_SNAKE_CASE__					=					4.0
          SCREAMING_SNAKE_CASE__					=					1e-6
          SCREAMING_SNAKE_CASE__					=					0.9_5
     else:
          raise ValueError(F'''Size {size} not supported'''       )

     # load image processor
     SCREAMING_SNAKE_CASE__					=					PoolFormerImageProcessor(crop_pct=_A       )

     # Prepare image
     SCREAMING_SNAKE_CASE__					=					prepare_img()
     SCREAMING_SNAKE_CASE__					=					image_processor(images=_A   ,	return_tensors='''pt'''       ).pixel_values

     logger.info(F'''Converting model {model_name}...'''       )

     # load original state dict
     SCREAMING_SNAKE_CASE__					=					torch.load(_A   ,	map_location=torch.device('''cpu'''       )       )

     # rename keys
     SCREAMING_SNAKE_CASE__					=					rename_keys(_A       )

     # create HuggingFace model and load state dict
     SCREAMING_SNAKE_CASE__					=					PoolFormerForImageClassification(_A       )
     model.load_state_dict(_A       )
     model.eval()

     # Define image processor
     SCREAMING_SNAKE_CASE__					=					PoolFormerImageProcessor(crop_pct=_A       )
     SCREAMING_SNAKE_CASE__					=					image_processor(images=prepare_img()   ,	return_tensors='''pt'''       ).pixel_values

     # forward pass
     SCREAMING_SNAKE_CASE__					=					model(_A       )
     SCREAMING_SNAKE_CASE__					=					outputs.logits

     # define expected logit slices for different models
     if size == "s12":
          SCREAMING_SNAKE_CASE__					=					torch.tensor([-0.3_0_4_5, -0.6_7_5_8, -0.4_8_6_9]       )
     elif size == "s24":
          SCREAMING_SNAKE_CASE__					=					torch.tensor([0.4_4_0_2, -0.1_3_7_4, -0.8_0_4_5]       )
     elif size == "s36":
          SCREAMING_SNAKE_CASE__					=					torch.tensor([-0.6_0_8_0, -0.5_1_3_3, -0.5_8_9_8]       )
     elif size == "m36":
          SCREAMING_SNAKE_CASE__					=					torch.tensor([0.3_9_5_2, 0.2_2_6_3, -1.2_6_6_8]       )
     elif size == "m48":
          SCREAMING_SNAKE_CASE__					=					torch.tensor([0.1_1_6_7, -0.0_6_5_6, -0.3_4_2_3]       )
     else:
          raise ValueError(F'''Size {size} not supported'''       )

     # verify logits
     assert logits.shape == expected_shape
     assert torch.allclose(logits[0, :3]   ,	_A   ,	atol=1e-2       )

     # finally, save model and image processor
     logger.info(F'''Saving PyTorch model and image processor to {pytorch_dump_folder_path}...'''       )
     Path(_A       ).mkdir(exist_ok=_A       )
     model.save_pretrained(_A       )
     print(F'''Saving image processor to {pytorch_dump_folder_path}'''       )
     image_processor.save_pretrained(_A       )


if __name__ == "__main__":
       _SCREAMING_SNAKE_CASE  :						Union[str, Any]         =    argparse.ArgumentParser()

       parser.add_argument(
           '''--model_name''',
           default='''poolformer_s12''',
           type=str,
           help='''Name of the model you\'d like to convert.''',
       )
       parser.add_argument(
           '''--checkpoint_path''', default=None, type=str, help='''Path to the original PyTorch checkpoint (.pth file).'''
       )
       parser.add_argument(
           '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the folder to output PyTorch model.'''
       )
       _SCREAMING_SNAKE_CASE  :						Tuple         =    parser.parse_args()
       convert_poolformer_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path)


