

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        lowerCamelCase__       (    _a):
	SCREAMING_SNAKE_CASE      :					List[str]    =				ASTConfig()

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

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

	SCREAMING_SNAKE_CASE      :					List[str]    =				json.load(open(hf_hub_download(_a							,	_a							,	repo_type="dataset")							,	"r"))
	SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				{int(_a): v for k, v in idalabel.items()}
	SCREAMING_SNAKE_CASE      :					Optional[Any]    =				idalabel
	SCREAMING_SNAKE_CASE      :					Tuple    =				{v: k for k, v in idalabel.items()}

	return config

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

	return name

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

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

	return orig_state_dict

def        lowerCamelCase__       (    _a):
	SCREAMING_SNAKE_CASE      :					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(_a							,	_a)

@torch.no_grad()
def        lowerCamelCase__       (    _a							,	_a							,	_a=False):
	SCREAMING_SNAKE_CASE      :					int    =				get_audio_spectrogram_transformer_config(_a)

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

	# load 🤗 model
	SCREAMING_SNAKE_CASE      :					Dict    =				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      :					List[str]    =				-4.267_7393 if "speech-commands" not in model_name else -6.84_5978
	SCREAMING_SNAKE_CASE      :					Tuple    =				4.568_9974 if "speech-commands" not in model_name else 5.565_4526
	SCREAMING_SNAKE_CASE      :					List[Any]    =				1024 if "speech-commands" not in model_name else 128
	SCREAMING_SNAKE_CASE      :					Optional[int]    =				ASTFeatureExtractor(mean=_a							,	std=_a							,	max_length=_a)

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

		SCREAMING_SNAKE_CASE   ,SCREAMING_SNAKE_CASE      :					List[str]    =				torchaudio.load(_a)
		SCREAMING_SNAKE_CASE      :					Any    =				waveform.squeeze().numpy()

	SCREAMING_SNAKE_CASE      :					Optional[Any]    =				feature_extractor(_a							,	sampling_rate=16000							,	return_tensors="pt")

	# forward pass
	SCREAMING_SNAKE_CASE      :					Optional[Any]    =				model(**_a)
	SCREAMING_SNAKE_CASE      :					int    =				outputs.logits

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



"""simple docstring"""







import argparse
import glob
import logging
import os
from argparse import Namespace
from importlib import import_module

import numpy as np
import torch
from lightning_base import BaseTransformer, add_generic_args, generic_train
from seqeval.metrics import accuracy_score, fa_score, precision_score, recall_score
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader, TensorDataset
from utils_ner import TokenClassificationTask


lowerCAmelCase__      				=   logging.getLogger(__name__)


class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       a  :				Tuple   ="token-classification"

       def __init__(						self		,     snake_case__						):






         """simple docstring"""





         if type(snake_case__						) == dict:
           lowerCAmelCase     :     List[str]								=							Namespace(**snake_case__						)
         lowerCAmelCase     :     Optional[Any]								=							import_module("tasks"						)
         try:
           lowerCAmelCase     :     Dict								=							getattr(snake_case__		,     hparams.task_type						)
           lowerCAmelCase     :     TokenClassificationTask								=							token_classification_task_clazz()
         except AttributeError:
           raise ValueError(
               f"""Task {hparams.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. """
               f"""Available tasks classes are: {TokenClassificationTask.__subclasses__()}"""						)
         lowerCAmelCase     :     Dict								=							self.token_classification_task.get_labels(hparams.labels						)
         lowerCAmelCase     :     str								=							CrossEntropyLoss().ignore_index
         super().__init__(snake_case__		,     len(self.labels						)		,     self.mode						)

       def        lowercase__							(						self		,     **snake_case__						):






         """simple docstring"""





         return self.model(**snake_case__						)

       def        lowercase__							(						self		,     snake_case__		,     snake_case__						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							{"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}
         if self.config.model_type != "distilbert":
           lowerCAmelCase     :     Optional[int]								=							(
               batch[2] if self.config.model_type in ["bert", "xlnet"] else None
           )  # XLM and RoBERTa don"t use token_type_ids

         lowerCAmelCase     :     str								=							self(**snake_case__						)
         lowerCAmelCase     :     Tuple								=							outputs[0]
         # tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
         return {"loss": loss}

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							self.hparams
         for mode in ["train", "dev", "test"]:
           lowerCAmelCase     :     Optional[int]								=							self._feature_file(snake_case__						)
           if os.path.exists(snake_case__						) and not args.overwrite_cache:
             logger.info("Loading features from cached file %s"		,     snake_case__						)
             lowerCAmelCase     :     Dict								=							torch.load(snake_case__						)
           else:
             logger.info("Creating features from dataset file at %s"		,     args.data_dir						)
             lowerCAmelCase     :     Optional[int]								=							self.token_classification_task.read_examples_from_file(args.data_dir		,     snake_case__						)
             lowerCAmelCase     :     Optional[int]								=							self.token_classification_task.convert_examples_to_features(
                 snake_case__		,     self.labels		,     args.max_seq_length		,     self.tokenizer		,     cls_token_at_end=bool(self.config.model_type in ["xlnet"]						)		,     cls_token=self.tokenizer.cls_token		,     cls_token_segment_id=2 if self.config.model_type in ["xlnet"] else 0		,     sep_token=self.tokenizer.sep_token		,     sep_token_extra=snake_case__		,     pad_on_left=bool(self.config.model_type in ["xlnet"]						)		,     pad_token=self.tokenizer.pad_token_id		,     pad_token_segment_id=self.tokenizer.pad_token_type_id		,     pad_token_label_id=self.pad_token_label_id		,     )
             logger.info("Saving features into cached file %s"		,     snake_case__						)
             torch.save(snake_case__		,     snake_case__						)

       def        lowercase__							(						self		,     snake_case__		,     snake_case__		,     snake_case__ = False						):






         """simple docstring"""





         lowerCAmelCase     :     Dict								=							self._feature_file(snake_case__						)
         logger.info("Loading features from cached file %s"		,     snake_case__						)
         lowerCAmelCase     :     Optional[Any]								=							torch.load(snake_case__						)
         lowerCAmelCase     :     Optional[int]								=							torch.tensor([f.input_ids for f in features]		,     dtype=torch.long						)
         lowerCAmelCase     :     int								=							torch.tensor([f.attention_mask for f in features]		,     dtype=torch.long						)
         if features[0].token_type_ids is not None:
           lowerCAmelCase     :     Tuple								=							torch.tensor([f.token_type_ids for f in features]		,     dtype=torch.long						)
         else:
           lowerCAmelCase     :     Optional[Any]								=							torch.tensor([0 for f in features]		,     dtype=torch.long						)
           # HACK(we will not use this anymore soon)
         lowerCAmelCase     :     Optional[int]								=							torch.tensor([f.label_ids for f in features]		,     dtype=torch.long						)
         return DataLoader(
             TensorDataset(snake_case__		,     snake_case__		,     snake_case__		,     snake_case__						)		,     batch_size=snake_case__						)

       def        lowercase__							(						self		,     snake_case__		,     snake_case__						):






         """simple docstring"""





         """Compute validation""" ""
         lowerCAmelCase     :     str								=							{"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}
         if self.config.model_type != "distilbert":
           lowerCAmelCase     :     List[Any]								=							(
               batch[2] if self.config.model_type in ["bert", "xlnet"] else None
           )  # XLM and RoBERTa don"t use token_type_ids
         lowerCAmelCase     :     Any								=							self(**snake_case__						)
         lowerCAmelCase						,  lowerCAmelCase     :     List[str]								=							outputs[:2]
         lowerCAmelCase     :     Dict								=							logits.detach().cpu().numpy()
         lowerCAmelCase     :     Optional[Any]								=							inputs["labels"].detach().cpu().numpy()
         return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}

       def        lowercase__							(						self		,     snake_case__						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							torch.stack([x["val_loss"] for x in outputs]						).mean()
         lowerCAmelCase     :     Optional[Any]								=							np.concatenate([x["pred"] for x in outputs]		,     axis=0						)
         lowerCAmelCase     :     List[Any]								=							np.argmax(snake_case__		,     axis=2						)
         lowerCAmelCase     :     Optional[int]								=							np.concatenate([x["target"] for x in outputs]		,     axis=0						)

         lowerCAmelCase     :     str								=							dict(enumerate(self.labels						)						)
         lowerCAmelCase     :     Optional[int]								=							[[] for _ in range(out_label_ids.shape[0]						)]
         lowerCAmelCase     :     Dict								=							[[] for _ in range(out_label_ids.shape[0]						)]

         for i in range(out_label_ids.shape[0]						):
           for j in range(out_label_ids.shape[1]						):
             if out_label_ids[i, j] != self.pad_token_label_id:
               out_label_list[i].append(label_map[out_label_ids[i][j]]						)
               preds_list[i].append(label_map[preds[i][j]]						)

         lowerCAmelCase     :     Dict								=							{
             "val_loss": val_loss_mean,
             "accuracy_score": accuracy_score(snake_case__		,     snake_case__						),
             "precision": precision_score(snake_case__		,     snake_case__						),
             "recall": recall_score(snake_case__		,     snake_case__						),
             "f1": fa_score(snake_case__		,     snake_case__						),
         }

         lowerCAmelCase     :     int								=							dict(results.items()						)
         lowerCAmelCase     :     Tuple								=							results
         return ret, preds_list, out_label_list

       def        lowercase__							(						self		,     snake_case__						):






         """simple docstring"""





         lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase     :     Dict								=							self._eval_end(snake_case__						)
         lowerCAmelCase     :     int								=							ret["log"]
         return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}

       def        lowercase__							(						self		,     snake_case__						):






         """simple docstring"""





         lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase     :     int								=							self._eval_end(snake_case__						)

         # Converting to the dict required by pl
         # https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
         # pytorch_lightning/trainer/logging.py#L139
         lowerCAmelCase     :     str								=							ret["log"]
         # `val_loss` is the key returned by `self._eval_end()` but actually refers to `test_loss`
         return {"avg_test_loss": logs["val_loss"], "log": logs, "progress_bar": logs}







       @staticmethod
       def        lowercase__							(						snake_case__		,     snake_case__						):






         """simple docstring"""





         BaseTransformer.add_model_specific_args(snake_case__		,     snake_case__						)
         parser.add_argument(
             "--task_type"		,     default="NER"		,     type=snake_case__		,     help="Task type to fine tune in training (e.g. NER, POS, etc)"						)
         parser.add_argument(
             "--max_seq_length"		,     default=128		,     type=snake_case__		,     help=(
                 "The maximum total input sequence length after tokenization. Sequences longer "
                 "than this will be truncated, sequences shorter will be padded."
             )		,     )

         parser.add_argument(
             "--labels"		,     default=""		,     type=snake_case__		,     help="Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."		,     )
         parser.add_argument(
             "--gpus"		,     default=0		,     type=snake_case__		,     help="The number of GPUs allocated for this, it is by default 0 meaning none"		,     )

         parser.add_argument(
             "--overwrite_cache"		,     action="store_true"		,     help="Overwrite the cached training and evaluation sets"						)

         return parser


if __name__ == "__main__":
     lowerCAmelCase__      				=   argparse.ArgumentParser()
     add_generic_args(parser, os.getcwd())
     lowerCAmelCase__      				=   NERTransformer.add_model_specific_args(parser, os.getcwd())
     lowerCAmelCase__      				=   parser.parse_args()
     lowerCAmelCase__      				=   NERTransformer(args)
     lowerCAmelCase__      				=   generic_train(model, args)

     if args.do_predict:
          # See https://github.com/huggingface/transformers/issues/3159
          # pl use this default format to create a checkpoint:
          # https://github.com/PyTorchLightning/pytorch-lightning/blob/master\
          # /pytorch_lightning/callbacks/model_checkpoint.py#L322
          lowerCAmelCase__      				=   sorted(glob.glob(os.path.join(args.output_dir, '''checkpoint-epoch=*.ckpt'''), recursive=True))
          lowerCAmelCase__      				=   model.load_from_checkpoint(checkpoints[-1])
          trainer.test(model)





