
"""simple docstring"""





import itertools
import json
import linecache
import os
import pickle
import re
import socket
import string
from collections import Counter
from logging import getLogger
from pathlib import Path
from typing import Callable, Dict, Iterable, List

import git
import torch
from torch.utils.data import Dataset

from transformers import BartTokenizer, RagTokenizer, TaTokenizer
def 				A_  (__a					, __a					, __a					, __a					, __a=True					, __a="pt"					):

					'''simple docstring'''





					A_								=    {"add_prefix_space": True} if isinstance(__a					, __a					) and not line.startswith(" "					) else {}
					A_								=    padding_side
					return tokenizer(
					    [line]					, max_length=__a					, padding="max_length" if pad_to_max_length else None					, truncation=__a					, return_tensors=__a					, add_special_tokens=__a					, **__a					, )
def 				A_  (__a					, __a					, __a=None					, ):

					'''simple docstring'''





					A_								=    input_ids.ne(__a					).any(dim=0					)
					if attention_mask is None:
										return input_ids[:, keep_column_mask]
					else:
										return (input_ids[:, keep_column_mask], attention_mask[:, keep_column_mask])



class __lowerCAmelCase					(       _lowercase   ):




			"""simple docstring"""







			def __init__(					self    :       Optional[Any]							, _snake_case    :       List[Any]							, _snake_case    :       List[Any]							, _snake_case    :       int							, _snake_case    :       Tuple							, _snake_case    :       Dict="train"							, _snake_case    :       List[Any]=None							, _snake_case    :       Optional[int]=None							, _snake_case    :       Optional[int]=None							, _snake_case    :       Any=""							, )   ->			List[str]:







								"""simple docstring"""
								super().__init__()
								A_								=    Path(_snake_case						).joinpath(type_path + ".source"						)
								A_								=    Path(_snake_case						).joinpath(type_path + ".target"						)
								A_								=    self.get_char_lens(self.src_file						)
								A_								=    max_source_length
								A_								=    max_target_length
								assert min(self.src_lens						) > 0, F'found empty line in {self.src_file}'
								A_								=    tokenizer
								A_								=    prefix
								if n_obs is not None:
													A_								=    self.src_lens[:n_obs]
								A_								=    src_lang
								A_								=    tgt_lang




			def __len__(					self    :       Optional[int]						)   ->			List[str]:







								"""simple docstring"""
								return len(self.src_lens						)




			def __getitem__(					self    :       str							, _snake_case    :       Optional[int]						)   ->			Dict[str, torch.Tensor]:







								"""simple docstring"""
								A_								=    index + 1  # linecache starts at 1
								A_								=    self.prefix + linecache.getline(str(self.src_file						)							, _snake_case						).rstrip("\n"						)
								A_								=    linecache.getline(str(self.tgt_file						)							, _snake_case						).rstrip("\n"						)
								assert source_line, F'empty source line for index {index}'
								assert tgt_line, F'empty tgt line for index {index}'

								# Need to add eos token manually for T5
								if isinstance(self.tokenizer							, _snake_case						):
													source_line += self.tokenizer.eos_token
													tgt_line += self.tokenizer.eos_token

								# Pad source and target to the right
								A_								=    (
								    self.tokenizer.question_encoder if isinstance(self.tokenizer							, _snake_case						) else self.tokenizer
								)
								A_								=    self.tokenizer.generator if isinstance(self.tokenizer							, _snake_case						) else self.tokenizer

								A_								=    encode_line(_snake_case							, _snake_case							, self.max_source_length							, "right"						)
								A_								=    encode_line(_snake_case							, _snake_case							, self.max_target_length							, "right"						)

								A_								=    source_inputs["input_ids"].squeeze()
								A_								=    target_inputs["input_ids"].squeeze()
								A_								=    source_inputs["attention_mask"].squeeze()
								return {
								    "input_ids": source_ids,
								    "attention_mask": src_mask,
								    "decoder_input_ids": target_ids,
								}




			@staticmethod
			def   lowerCamelCase__  (					_snake_case    :       List[Any]						)   ->			str:







								"""simple docstring"""
								return [len(_snake_case						) for x in Path(_snake_case						).open().readlines()]






			def   lowerCamelCase__  (					self    :       Any							, _snake_case    :       List[str]						)   ->			Dict[str, torch.Tensor]:







								"""simple docstring"""
								A_								=    torch.stack([x["input_ids"] for x in batch]						)
								A_								=    torch.stack([x["attention_mask"] for x in batch]						)
								A_								=    torch.stack([x["decoder_input_ids"] for x in batch]						)
								A_								=    (
								    self.tokenizer.generator.pad_token_id
								    if isinstance(self.tokenizer							, _snake_case						)
								    else self.tokenizer.pad_token_id
								)
								A_								=    (
								    self.tokenizer.question_encoder.pad_token_id
								    if isinstance(self.tokenizer							, _snake_case						)
								    else self.tokenizer.pad_token_id
								)
								A_								=    trim_batch(_snake_case							, _snake_case						)
								A_			,     A_								=    trim_batch(_snake_case							, _snake_case							, attention_mask=_snake_case						)
								A_								=    {
								    "input_ids": source_ids,
								    "attention_mask": source_mask,
								    "decoder_input_ids": y,
								}
								return batch





UpperCamelCase_       :    Any						      =			getLogger(__name__)
def 				A_  (__a					):

					'''simple docstring'''





					return list(itertools.chain.from_iterable(__a					)					)
def 				A_  (__a					):

					'''simple docstring'''





					A_								=    get_git_info()
					save_json(__a					, os.path.join(__a					, "git_log.json"					)					)
def 				A_  (__a					, __a					, __a=4					, **__a					):

					'''simple docstring'''





					with open(__a					, "w"					) as f:
										json.dump(__a					, __a					, indent=__a					, **__a					)
def 				A_  (__a					):

					'''simple docstring'''





					with open(__a					) as f:
										return json.load(__a					)
def 				A_  ():

					'''simple docstring'''





					A_								=    git.Repo(search_parent_directories=__a					)
					A_								=    {
					    "repo_id": str(__a					),
					    "repo_sha": str(repo.head.object.hexsha					),
					    "repo_branch": str(repo.active_branch					),
					    "hostname": str(socket.gethostname()					),
					}
					return repo_infos
def 				A_  (__a					, __a					):

					'''simple docstring'''





					return list(map(__a					, __a					)					)
def 				A_  (__a					, __a					):

					'''simple docstring'''





					with open(__a					, "wb"					) as f:
										return pickle.dump(__a					, __a					)
def 				A_  (__a					):

					'''simple docstring'''





					def remove_articles(__a					):
										return re.sub(R"\b(a|an|the)\b"					, " "					, __a					)

					def white_space_fix(__a					):
										return " ".join(text.split()					)

					def remove_punc(__a					):
										A_								=    set(string.punctuation					)
										return "".join(ch for ch in text if ch not in exclude					)

					def lower(__a					):
										return text.lower()

					return white_space_fix(remove_articles(remove_punc(lower(__a					)					)					)					)
def 				A_  (__a					, __a					):

					'''simple docstring'''





					A_								=    normalize_answer(__a					).split()
					A_								=    normalize_answer(__a					).split()
					A_								=    Counter(__a					) & Counter(__a					)
					A_								=    sum(common.values()					)
					if num_same == 0:
										return 0
					A_								=    1.0 * num_same / len(__a					)
					A_								=    1.0 * num_same / len(__a					)
					A_								=    (2 * precision * recall) / (precision + recall)
					return fa
def 				A_  (__a					, __a					):

					'''simple docstring'''





					return normalize_answer(__a					) == normalize_answer(__a					)
def 				A_  (__a					, __a					):

					'''simple docstring'''





					assert len(__a					) == len(__a					)
					A_								=    0
					for hypo, pred in zip(__a					, __a					):
										em += exact_match_score(__a					, __a					)
					if len(__a					) > 0:
										em /= len(__a					)
					return {"em": em}
def 				A_  (__a					):

					'''simple docstring'''





					return model_prefix.startswith("rag"					)




def 				A_  (__a					, __a					, __a					):

					'''simple docstring'''





					A_								=    {p: p for p in extra_params}
					# T5 models don't have `dropout` param, they have `dropout_rate` instead
					A_								=    "dropout_rate"
					for p in extra_params:
										if getattr(__a					, __a					, __a					):
															if not hasattr(__a					, __a					) and not hasattr(__a					, equivalent_param[p]					):
																				logger.info("config doesn't have a `{}` attribute".format(__a					)					)
																				delattr(__a					, __a					)
																				continue
															A_								=    p if hasattr(__a					, __a					) else equivalent_param[p]
															setattr(__a					, __a					, getattr(__a					, __a					)					)
															delattr(__a					, __a					)
					return hparams, config






"""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()
UpperCamelCase_       :    Any						      =			logging.get_logger(__name__)
def 				A_  (__a					):

					'''simple docstring'''





					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(__a					, __a					, repo_type="dataset"					)					, "r"					)					)
					A_								=    {int(__a					): v for k, v in idalabel.items()}
					A_								=    idalabel
					A_								=    {v: k for k, v in idalabel.items()}

					return config
def 				A_  (__a					):

					'''simple docstring'''





					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 				A_  (__a					, __a					):

					'''simple docstring'''





					for key in orig_state_dict.copy().keys():
										A_								=    orig_state_dict.pop(__a					)

										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 				A_  (__a					):

					'''simple docstring'''





					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(__a					, __a					)
@torch.no_grad()
def 				A_  (__a					, __a					, __a=False					):

					'''simple docstring'''





					A_								=    get_audio_spectrogram_transformer_config(__a					)

					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(__a					, map_location="cpu"					)
					# remove some keys
					remove_keys(__a					)
					# rename some keys
					A_								=    convert_state_dict(__a					, __a					)

					# load 🤗 model
					A_								=    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
					A_								=    -4.2677393 if "speech-commands" not in model_name else -6.845978
					A_								=    4.5689974 if "speech-commands" not in model_name else 5.5654526
					A_								=    1024 if "speech-commands" not in model_name else 128
					A_								=    ASTFeatureExtractor(mean=__a					, std=__a					, max_length=__a					)

					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(__a					)
										A_								=    waveform.squeeze().numpy()

					A_								=    feature_extractor(__a					, sampling_rate=1_6000					, return_tensors="pt"					)

					# forward pass
					A_								=    model(**__a					)
					A_								=    outputs.logits

					if model_name == "ast-finetuned-audioset-10-10-0.4593":
										A_								=    torch.tensor([-0.8760, -7.0042, -8.6602]					)
					elif model_name == "ast-finetuned-audioset-10-10-0.450":
										A_								=    torch.tensor([-1.1986, -7.0903, -8.2718]					)
					elif model_name == "ast-finetuned-audioset-10-10-0.448":
										A_								=    torch.tensor([-2.6128, -8.0080, -9.4344]					)
					elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
										A_								=    torch.tensor([-1.5080, -7.4534, -8.8917]					)
					elif model_name == "ast-finetuned-audioset-12-12-0.447":
										A_								=    torch.tensor([-0.5050, -6.5833, -8.0843]					)
					elif model_name == "ast-finetuned-audioset-14-14-0.443":
										A_								=    torch.tensor([-0.3826, -7.0336, -8.2413]					)
					elif model_name == "ast-finetuned-audioset-16-16-0.442":
										A_								=    torch.tensor([-1.2113, -6.9101, -8.3470]					)
					elif model_name == "ast-finetuned-speech-commands-v2":
										A_								=    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__":
		UpperCamelCase_       :    Dict						      =			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_       :    List[str]						      =			parser.parse_args()
		convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




