





'''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()
lowerCAmelCase_					:    Union[str, Any]      	=			logging.get_logger(__name__)






def  _lowerCamelCase				( lowercase						:      List[Any]						)	->	str:
				_a								= ASTConfig()

				if "10-10" in model_name:
								pass
				elif "speech-commands" in model_name:
								_a								= 128
				elif "12-12" in model_name:
								_a								= 12
								_a								= 12
				elif "14-14" in model_name:
								_a								= 14
								_a								= 14
				elif "16-16" in model_name:
								_a								= 16
								_a								= 16
				else:
								raise ValueError("Model not supported"						)

				_a								= "huggingface/label-files"
				if "speech-commands" in model_name:
								_a								= 35
								_a								= "speech-commands-v2-id2label.json"
				else:
								_a								= 527
								_a								= "audioset-id2label.json"

				_a								= json.load(open(hf_hub_download(lowercase    ,							lowercase    ,							repo_type="dataset"						)    ,							"r"						)						)
				_a								= {int(lowercase						): v for k, v in idalabel.items()}
				_a								= idalabel
				_a								= {v: k for k, v in idalabel.items()}

				return config






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

				return name






def  _lowerCamelCase				( lowercase						:      Any    ,							lowercase						:      Tuple						)	->	List[str]:
				for key in orig_state_dict.copy().keys():
								_a								= orig_state_dict.pop(lowercase						)

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

				return orig_state_dict






def  _lowerCamelCase				( lowercase						:      Optional[int]						)	->	Optional[int]:
				_a								= [
				    "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(lowercase    ,							lowercase						)






@torch.no_grad()
def  _lowerCamelCase				( lowercase						:      Union[str, Any]    ,							lowercase						:      List[Any]    ,							lowercase						:      Dict=False						)	->	int:
				_a								= get_audio_spectrogram_transformer_config(lowercase						)

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

				# load 🤗 model
				_a								= ASTForAudioClassification(lowercase						)
				model.eval()

				model.load_state_dict(lowercase						)

				# verify outputs on dummy input
				# source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
				_a								= -4.2_67_73_93 if "speech-commands" not in model_name else -6.84_59_78
				_a								= 4.5_68_99_74 if "speech-commands" not in model_name else 5.5_65_45_26
				_a								= 1024 if "speech-commands" not in model_name else 128
				_a								= ASTFeatureExtractor(mean=lowercase    ,							std=lowercase    ,							max_length=lowercase						)

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

								_a	,					_a								= torchaudio.load(lowercase						)
								_a								= waveform.squeeze().numpy()

				_a								= feature_extractor(lowercase    ,							sampling_rate=1_6000    ,							return_tensors="pt"						)

				# forward pass
				_a								= model(**lowercase						)
				_a								= outputs.logits

				if model_name == "ast-finetuned-audioset-10-10-0.4593":
								_a								= torch.tensor([-0.87_60, -7.00_42, -8.66_02]						)
				elif model_name == "ast-finetuned-audioset-10-10-0.450":
								_a								= torch.tensor([-1.19_86, -7.09_03, -8.27_18]						)
				elif model_name == "ast-finetuned-audioset-10-10-0.448":
								_a								= torch.tensor([-2.61_28, -8.00_80, -9.43_44]						)
				elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
								_a								= torch.tensor([-1.50_80, -7.45_34, -8.89_17]						)
				elif model_name == "ast-finetuned-audioset-12-12-0.447":
								_a								= torch.tensor([-0.50_50, -6.58_33, -8.08_43]						)
				elif model_name == "ast-finetuned-audioset-14-14-0.443":
								_a								= torch.tensor([-0.38_26, -7.03_36, -8.24_13]						)
				elif model_name == "ast-finetuned-audioset-16-16-0.442":
								_a								= torch.tensor([-1.21_13, -6.91_01, -8.34_70]						)
				elif model_name == "ast-finetuned-speech-commands-v2":
								_a								= torch.tensor([6.15_89, -8.05_66, -8.79_84]						)
				else:
								raise ValueError("Unknown model name"						)
				if not torch.allclose(logits[0, :3]    ,							lowercase    ,							atol=1E-4						):
								raise ValueError("Logits don't match"						)
				print("Looks ok!"						)

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

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

					lowerCAmelCase_					:    Dict      	=			parser.parse_args()
					convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase__   	=							logging.get_logger(__name__)

UpperCAmelCase__   	=							{
    '''microsoft/wavlm-base''': '''https://huggingface.co/microsoft/wavlm-base/resolve/main/config.json''',
    # See all WavLM models at https://huggingface.co/models?filter=wavlm
}







class       lowerCamelCase__   (       lowerCAmelCase):
     SCREAMING_SNAKE_CASE__      			=			'''wavlm'''





     def __init__(self		,     UpperCAmelCase=3_2		,     UpperCAmelCase=7_6_8		,     UpperCAmelCase=1_2		,     UpperCAmelCase=1_2		,     UpperCAmelCase=3_0_7_2		,     UpperCAmelCase="gelu"		,     UpperCAmelCase=0.1		,     UpperCAmelCase=0.1		,     UpperCAmelCase=0.1		,     UpperCAmelCase=0.0		,     UpperCAmelCase=0.1		,     UpperCAmelCase=0.1		,     UpperCAmelCase=0.02		,     UpperCAmelCase=1e-5		,     UpperCAmelCase="group"		,     UpperCAmelCase="gelu"		,     UpperCAmelCase=(5_1_2, 5_1_2, 5_1_2, 5_1_2, 5_1_2, 5_1_2, 5_1_2)		,     UpperCAmelCase=(5, 2, 2, 2, 2, 2, 2)		,     UpperCAmelCase=(1_0, 3, 3, 3, 3, 2, 2)		,     UpperCAmelCase=False		,     UpperCAmelCase=1_2_8		,     UpperCAmelCase=1_6		,     UpperCAmelCase=3_2_0		,     UpperCAmelCase=8_0_0		,     UpperCAmelCase=False		,     UpperCAmelCase=True		,     UpperCAmelCase=0.05		,     UpperCAmelCase=1_0		,     UpperCAmelCase=2		,     UpperCAmelCase=0.0		,     UpperCAmelCase=1_0		,     UpperCAmelCase=3_2_0		,     UpperCAmelCase=2		,     UpperCAmelCase=0.1		,     UpperCAmelCase=1_0_0		,     UpperCAmelCase=2_5_6		,     UpperCAmelCase=2_5_6		,     UpperCAmelCase=0.1		,     UpperCAmelCase="mean"		,     UpperCAmelCase=False		,     UpperCAmelCase=False		,     UpperCAmelCase=2_5_6		,     UpperCAmelCase=(5_1_2, 5_1_2, 5_1_2, 5_1_2, 1_5_0_0)		,     UpperCAmelCase=(5, 3, 3, 1, 1)		,     UpperCAmelCase=(1, 2, 3, 1, 1)		,     UpperCAmelCase=5_1_2		,     UpperCAmelCase=8_0		,     UpperCAmelCase=0		,     UpperCAmelCase=1		,     UpperCAmelCase=2		,     UpperCAmelCase=False		,     UpperCAmelCase=3		,     UpperCAmelCase=2		,     UpperCAmelCase=3		,     UpperCAmelCase=None		,     **UpperCAmelCase		,     )     ->   Optional[Any]:
      super().__init__(**UpperCAmelCase		,     pad_token_id=UpperCAmelCase		,     bos_token_id=UpperCAmelCase		,     eos_token_id=UpperCAmelCase       )
      _lowercase							      =hidden_size
      _lowercase							      =feat_extract_norm
      _lowercase							      =feat_extract_activation
      _lowercase							      =list(UpperCAmelCase       )
      _lowercase							      =list(UpperCAmelCase       )
      _lowercase							      =list(UpperCAmelCase       )
      _lowercase							      =conv_bias
      _lowercase							      =num_buckets
      _lowercase							      =max_bucket_distance
      _lowercase							      =num_conv_pos_embeddings
      _lowercase							      =num_conv_pos_embedding_groups
      _lowercase							      =len(self.conv_dim       )
      _lowercase							      =num_hidden_layers
      _lowercase							      =intermediate_size
      _lowercase							      =hidden_act
      _lowercase							      =num_attention_heads
      _lowercase							      =hidden_dropout
      _lowercase							      =attention_dropout
      _lowercase							      =activation_dropout
      _lowercase							      =feat_proj_dropout
      _lowercase							      =final_dropout
      _lowercase							      =layerdrop
      _lowercase							      =layer_norm_eps
      _lowercase							      =initializer_range
      _lowercase							      =num_ctc_classes
      _lowercase							      =vocab_size
      _lowercase							      =do_stable_layer_norm
      _lowercase							      =use_weighted_layer_sum
      _lowercase							      =classifier_proj_size

      if (
          (len(self.conv_stride       ) != self.num_feat_extract_layers)
          or (len(self.conv_kernel       ) != self.num_feat_extract_layers)
          or (len(self.conv_dim       ) != self.num_feat_extract_layers)
      ):
       raise ValueError(
           '''Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =='''
           ''' `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ='''
           f" {len(self.conv_dim       )}`, `len(config.conv_stride) = {len(self.conv_stride       )}`,"
           f" `len(config.conv_kernel) = {len(self.conv_kernel       )}`."       )

      # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
      _lowercase							      =apply_spec_augment
      _lowercase							      =mask_time_prob
      _lowercase							      =mask_time_length
      _lowercase							      =mask_time_min_masks
      _lowercase							      =mask_feature_prob
      _lowercase							      =mask_feature_length

      # parameters for pretraining with codevector quantized representations
      _lowercase							      =num_codevectors_per_group
      _lowercase							      =num_codevector_groups
      _lowercase							      =contrastive_logits_temperature
      _lowercase							      =num_negatives
      _lowercase							      =codevector_dim
      _lowercase							      =proj_codevector_dim
      _lowercase							      =diversity_loss_weight

      # ctc loss
      _lowercase							      =ctc_loss_reduction
      _lowercase							      =ctc_zero_infinity

      # adapter
      _lowercase							      =add_adapter
      _lowercase							      =adapter_kernel_size
      _lowercase							      =adapter_stride
      _lowercase							      =num_adapter_layers
      _lowercase							      =output_hidden_size or hidden_size

      # SequenceClassification-specific parameter. Feel free to ignore for other classes.
      _lowercase							      =classifier_proj_size

      # XVector-specific parameters. Feel free to ignore for other classes.
      _lowercase							      =list(UpperCAmelCase       )
      _lowercase							      =list(UpperCAmelCase       )
      _lowercase							      =list(UpperCAmelCase       )
      _lowercase							      =xvector_output_dim





     @property
     def      __A					(self       )     ->   int:
      return functools.reduce(operator.mul		,     self.conv_stride		,     1       )

