






import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ASTConfig
from transformers.testing_utils import require_torch, require_torchaudio, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_torchaudio_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch
					from torch import nn

					from transformers import ASTForAudioClassification, ASTModel
					from transformers.models.audio_spectrogram_transformer.modeling_audio_spectrogram_transformer import (
					    AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
					)


if is_torchaudio_available():
					import torchaudio

					from transformers import ASTFeatureExtractor


class     __SCREAMING_SNAKE_CASE					:






						def __init__(	self     ,						_UpperCamelCase     ,						_UpperCamelCase=13     ,						_UpperCamelCase=2     ,						_UpperCamelCase=24     ,						_UpperCamelCase=16     ,						_UpperCamelCase=True     ,						_UpperCamelCase=True     ,						_UpperCamelCase=32     ,						_UpperCamelCase=5     ,						_UpperCamelCase=4     ,						_UpperCamelCase=37     ,						_UpperCamelCase="gelu"     ,						_UpperCamelCase=0.1     ,						_UpperCamelCase=0.1     ,						_UpperCamelCase=10     ,						_UpperCamelCase=0.02     ,						_UpperCamelCase=None     ,						_UpperCamelCase=2     ,						_UpperCamelCase=2     ,						):






												"""simple docstring"""
												lowerCAmelCase__    						=						parent
												lowerCAmelCase__    						=						batch_size
												lowerCAmelCase__    						=						patch_size
												lowerCAmelCase__    						=						max_length
												lowerCAmelCase__    						=						num_mel_bins
												lowerCAmelCase__    						=						is_training
												lowerCAmelCase__    						=						use_labels
												lowerCAmelCase__    						=						hidden_size
												lowerCAmelCase__    						=						num_hidden_layers
												lowerCAmelCase__    						=						num_attention_heads
												lowerCAmelCase__    						=						intermediate_size
												lowerCAmelCase__    						=						hidden_act
												lowerCAmelCase__    						=						hidden_dropout_prob
												lowerCAmelCase__    						=						attention_probs_dropout_prob
												lowerCAmelCase__    						=						type_sequence_label_size
												lowerCAmelCase__    						=						initializer_range
												lowerCAmelCase__    						=						scope
												lowerCAmelCase__    						=						frequency_stride
												lowerCAmelCase__    						=						time_stride

												# in AST, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
												lowerCAmelCase__    						=						(self.num_mel_bins - self.patch_size) // self.frequency_stride + 1
												lowerCAmelCase__    						=						(self.max_length - self.patch_size) // self.time_stride + 1
												lowerCAmelCase__    						=						frequency_out_dimension * time_out_dimension
												lowerCAmelCase__    						=						num_patches + 2






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						floats_tensor([self.batch_size, self.max_length, self.num_mel_bins]   )

												lowerCAmelCase__    						=						None
												if self.use_labels:
																		lowerCAmelCase__    						=						ids_tensor([self.batch_size]     ,						self.type_sequence_label_size   )

												lowerCAmelCase__    						=						self.get_config()

												return config, input_values, labels






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												return ASTConfig(
												    patch_size=self.patch_size     ,						max_length=self.max_length     ,						num_mel_bins=self.num_mel_bins     ,						hidden_size=self.hidden_size     ,						num_hidden_layers=self.num_hidden_layers     ,						num_attention_heads=self.num_attention_heads     ,						intermediate_size=self.intermediate_size     ,						hidden_act=self.hidden_act     ,						hidden_dropout_prob=self.hidden_dropout_prob     ,						attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,						is_decoder=_SCREAMING_SNAKE_CASE     ,						initializer_range=self.initializer_range     ,						frequency_stride=self.frequency_stride     ,						time_stride=self.time_stride     ,						)






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						ASTModel(config=_SCREAMING_SNAKE_CASE   )
												model.to(_SCREAMING_SNAKE_CASE   )
												model.eval()
												lowerCAmelCase__    						=						model(_SCREAMING_SNAKE_CASE   )
												self.parent.assertEqual(result.last_hidden_state.shape     ,						(self.batch_size, self.seq_length, self.hidden_size)   )







						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.prepare_config_and_inputs()
												(
												    (
												    lowerCAmelCase__
												)						,       (
												    lowerCAmelCase__
												)						,       (
												    lowerCAmelCase__
												)						,       
												)    						=						config_and_inputs
												lowerCAmelCase__    						=						{'input_values': input_values}
												return config, inputs_dict




@require_torch
class     __SCREAMING_SNAKE_CASE					(   __lowercase		,					__lowercase		,					unittest.TestCase):
						_SCREAMING_SNAKE_CASE	:       Union[str, Any]   =					(
						    (
						        ASTModel,
						        ASTForAudioClassification,
						    )
						    if is_torch_available()
						    else ()
						)
						_SCREAMING_SNAKE_CASE	:       Tuple   =					(
						    {'''audio-classification''': ASTForAudioClassification, '''feature-extraction''': ASTModel}
						    if is_torch_available()
						    else {}
						)
						_SCREAMING_SNAKE_CASE	:       Any   =					False
						_SCREAMING_SNAKE_CASE	:       List[Any]   =					False
						_SCREAMING_SNAKE_CASE	:       List[str]   =					False
						_SCREAMING_SNAKE_CASE	:       List[Any]   =					False






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase   ):






												"""simple docstring"""






												if pipeline_test_casse_name == "AudioClassificationPipelineTests":
																		return True

												return False






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						ASTModelTester(self   )
												lowerCAmelCase__    						=						ConfigTester(self     ,						config_class=_SCREAMING_SNAKE_CASE     ,						has_text_modality=_SCREAMING_SNAKE_CASE     ,						hidden_size=37   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												self.config_tester.run_common_tests()






						@unittest.skip(reason='AST does not use inputs_embeds'   )
						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												pass






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__						,       lowerCAmelCase__    						=						self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																		lowerCAmelCase__    						=						model_class(_SCREAMING_SNAKE_CASE   )
																		self.assertIsInstance(model.get_input_embeddings()     ,						(nn.Module)   )
																		lowerCAmelCase__    						=						model.get_output_embeddings()
																		self.assertTrue(x is None or isinstance(_SCREAMING_SNAKE_CASE     ,						nn.Linear   )   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__						,       lowerCAmelCase__    						=						self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																		lowerCAmelCase__    						=						model_class(_SCREAMING_SNAKE_CASE   )
																		lowerCAmelCase__    						=						inspect.signature(model.forward   )
																		# signature.parameters is an OrderedDict => so arg_names order is deterministic
																		lowerCAmelCase__    						=						[*signature.parameters.keys()]

																		lowerCAmelCase__    						=						['input_values']
																		self.assertListEqual(arg_names[:1]     ,						_SCREAMING_SNAKE_CASE   )






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE   )







						@slow
						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""






												for model_name in AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		lowerCAmelCase__    						=						ASTModel.from_pretrained(_SCREAMING_SNAKE_CASE   )
																		self.assertIsNotNone(_SCREAMING_SNAKE_CASE   )







def 	_UpperCamelCase   (    )  ->   Union[str, Any]:

						"""simple docstring"""



						lowerCAmelCase__    						=						hf_hub_download(
						    repo_id='nielsr/audio-spectogram-transformer-checkpoint'    ,       filename='sample_audio.flac'    ,       repo_type='dataset'		)

						lowerCAmelCase__						,       lowerCAmelCase__    						=						torchaudio.load(__UpperCamelCase		)

						return audio, sampling_rate


@require_torch
@require_torchaudio
class     __SCREAMING_SNAKE_CASE					(   unittest.TestCase):






						@cached_property
						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												return (
												    ASTFeatureExtractor.from_pretrained('MIT/ast-finetuned-audioset-10-10-0.4593'   )
												    if is_torchaudio_available()
												    else None
												)







						@slow
						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.default_feature_extractor
												lowerCAmelCase__    						=						ASTForAudioClassification.from_pretrained('MIT/ast-finetuned-audioset-10-10-0.4593'   ).to(_SCREAMING_SNAKE_CASE   )

												lowerCAmelCase__    						=						self.default_feature_extractor
												lowerCAmelCase__						,       lowerCAmelCase__    						=						prepare_audio()
												lowerCAmelCase__    						=						audio.squeeze().numpy()
												lowerCAmelCase__    						=						feature_extractor(_SCREAMING_SNAKE_CASE     ,						sampling_rate=_SCREAMING_SNAKE_CASE     ,						return_tensors='pt'   ).to(_SCREAMING_SNAKE_CASE   )

												# forward pass
												with torch.no_grad():
																		lowerCAmelCase__    						=						model(**_SCREAMING_SNAKE_CASE   )

												# verify the logits
												lowerCAmelCase__    						=						torch.Size((1, 5_27)   )
												self.assertEqual(outputs.logits.shape     ,						_SCREAMING_SNAKE_CASE   )

												lowerCAmelCase__    						=						torch.tensor([-0.87_60, -7.00_42, -8.66_02]   ).to(_SCREAMING_SNAKE_CASE   )

												self.assertTrue(torch.allclose(outputs.logits[0, :3]     ,						_SCREAMING_SNAKE_CASE     ,						atol=1E-4   )   )










import logging
import os
import sys
from dataclasses import dataclass, field
from importlib import import_module
from typing import Dict, List, Optional, Tuple

import numpy as np
from seqeval.metrics import accuracy_score, fa_score, precision_score, recall_score
from torch import nn
from utils_ner import Split, TokenClassificationDataset, TokenClassificationTask

import transformers
from transformers import (
    AutoConfig,
    AutoModelForTokenClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    set_seed,
)
from transformers.trainer_utils import is_main_process


__snake_case    :    List[Any]         =     logging.getLogger(__name__)


@dataclass
class     __SCREAMING_SNAKE_CASE					:
						_SCREAMING_SNAKE_CASE	:       str   =					field(
						    metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models'''})
						_SCREAMING_SNAKE_CASE	:       Optional[str]   =					field(
						    default=__lowercase		,					metadata={'''help''': '''Pretrained config name or path if not the same as model_name'''})
						_SCREAMING_SNAKE_CASE	:       Optional[str]   =					field(
						    default='''NER'''		,					metadata={'''help''': '''Task type to fine tune in training (e.g. NER, POS, etc)'''})
						_SCREAMING_SNAKE_CASE	:       Optional[str]   =					field(
						    default=__lowercase		,					metadata={'''help''': '''Pretrained tokenizer name or path if not the same as model_name'''})
						_SCREAMING_SNAKE_CASE	:       bool   =					field(default=__lowercase		,					metadata={'''help''': '''Set this flag to use fast tokenization.'''})
						# If you want to tweak more attributes on your tokenizer, you should do it in a distinct script,
						# or just modify its tokenizer_config.json.
						_SCREAMING_SNAKE_CASE	:       Optional[str]   =					field(
						    default=__lowercase		,					metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co'''}		,					)




@dataclass
class     __SCREAMING_SNAKE_CASE					:
						_SCREAMING_SNAKE_CASE	:       str   =					field(
						    metadata={'''help''': '''The input data dir. Should contain the .txt files for a CoNLL-2003-formatted task.'''})
						_SCREAMING_SNAKE_CASE	:       Optional[str]   =					field(
						    default=__lowercase		,					metadata={'''help''': '''Path to a file containing all labels. If not specified, CoNLL-2003 labels are used.'''}		,					)
						_SCREAMING_SNAKE_CASE	:       int   =					field(
						    default=128		,					metadata={
						        '''help''': (
						            '''The maximum total input sequence length after tokenization. Sequences longer '''
						            '''than this will be truncated, sequences shorter will be padded.'''
						        )
						    }		,					)
						_SCREAMING_SNAKE_CASE	:       bool   =					field(
						    default=__lowercase		,					metadata={'''help''': '''Overwrite the cached training and evaluation sets'''})







def 	_UpperCamelCase   (    )  ->   List[str]:

						"""simple docstring"""



						lowerCAmelCase__    						=						HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)		)
						if len(sys.argv		) == 2 and sys.argv[1].endswith('.json'		):
												# If we pass only one argument to the script and it's the path to a json file,
												# let's parse it to get our arguments.
												lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__    						=						parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]		)		)
						else:
												lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__    						=						parser.parse_args_into_dataclasses()

						if (
						    os.path.exists(training_args.output_dir		)
						    and os.listdir(training_args.output_dir		)
						    and training_args.do_train
						    and not training_args.overwrite_output_dir
						):
												raise ValueError(
												    F"Output directory ({training_args.output_dir}) already exists and is not empty. Use"
												    ' --overwrite_output_dir to overcome.'		)

						lowerCAmelCase__    						=						import_module('tasks'		)
						try:
												lowerCAmelCase__    						=						getattr(UpperCamelCase_    ,       model_args.task_type		)
												lowerCAmelCase__    						=						token_classification_task_clazz()
						except AttributeError:
												raise ValueError(
												    F"Task {model_args.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. "
												    F"Available tasks classes are: {TokenClassificationTask.__subclasses__()}"		)

						# Setup logging
						logging.basicConfig(
						    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'    ,       datefmt='%m/%d/%Y %H:%M:%S'    ,       level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN    ,       )
						logger.warning(
						    'Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s'    ,       training_args.local_rank    ,       training_args.device    ,       training_args.n_gpu    ,       bool(training_args.local_rank != -1		)    ,       training_args.fpaa    ,       )
						# Set the verbosity to info of the Transformers logger (on main process only):
						if is_main_process(training_args.local_rank		):
												transformers.utils.logging.set_verbosity_info()
												transformers.utils.logging.enable_default_handler()
												transformers.utils.logging.enable_explicit_format()
						logger.info('Training/evaluation parameters %s'    ,       UpperCamelCase_		)

						# Set seed
						set_seed(training_args.seed		)

						# Prepare CONLL-2003 task
						lowerCAmelCase__    						=						token_classification_task.get_labels(data_args.labels		)
						lowerCAmelCase__    						=						dict(enumerate(UpperCamelCase_		)		)
						lowerCAmelCase__    						=						len(UpperCamelCase_		)

						# Load pretrained model and tokenizer
						#
						# Distributed training:
						# The .from_pretrained methods guarantee that only one local process can concurrently
						# download model & vocab.

						lowerCAmelCase__    						=						AutoConfig.from_pretrained(
						    model_args.config_name if model_args.config_name else model_args.model_name_or_path    ,       num_labels=UpperCamelCase_    ,       idalabel=UpperCamelCase_    ,       labelaid={label: i for i, label in enumerate(UpperCamelCase_		)}    ,       cache_dir=model_args.cache_dir    ,       )
						lowerCAmelCase__    						=						AutoTokenizer.from_pretrained(
						    model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path    ,       cache_dir=model_args.cache_dir    ,       use_fast=model_args.use_fast    ,       )
						lowerCAmelCase__    						=						AutoModelForTokenClassification.from_pretrained(
						    model_args.model_name_or_path    ,       from_tf=bool('.ckpt' in model_args.model_name_or_path		)    ,       config=UpperCamelCase_    ,       cache_dir=model_args.cache_dir    ,       )

						# Get datasets
						lowerCAmelCase__    						=						(
						    TokenClassificationDataset(
						        token_classification_task=UpperCamelCase_    ,       data_dir=data_args.data_dir    ,       tokenizer=UpperCamelCase_    ,       labels=UpperCamelCase_    ,       model_type=config.model_type    ,       max_seq_length=data_args.max_seq_length    ,       overwrite_cache=data_args.overwrite_cache    ,       mode=Split.train    ,       )
						    if training_args.do_train
						    else None
						)
						lowerCAmelCase__    						=						(
						    TokenClassificationDataset(
						        token_classification_task=UpperCamelCase_    ,       data_dir=data_args.data_dir    ,       tokenizer=UpperCamelCase_    ,       labels=UpperCamelCase_    ,       model_type=config.model_type    ,       max_seq_length=data_args.max_seq_length    ,       overwrite_cache=data_args.overwrite_cache    ,       mode=Split.dev    ,       )
						    if training_args.do_eval
						    else None
						)

						def align_predictions(UpperCamelCase_       :					np.ndarray    ,       UpperCamelCase_       :					np.ndarray		) -> Tuple[List[int], List[int]]:
												lowerCAmelCase__    						=						np.argmax(UpperCamelCase_    ,       axis=2		)

												lowerCAmelCase__						,       lowerCAmelCase__    						=						preds.shape

												lowerCAmelCase__    						=						[[] for _ in range(UpperCamelCase_		)]
												lowerCAmelCase__    						=						[[] for _ in range(UpperCamelCase_		)]

												for i in range(UpperCamelCase_		):
																		for j in range(UpperCamelCase_		):
																								if label_ids[i, j] != nn.CrossEntropyLoss().ignore_index:
																														out_label_list[i].append(label_map[label_ids[i][j]]		)
																														preds_list[i].append(label_map[preds[i][j]]		)

												return preds_list, out_label_list

						def compute_metrics(UpperCamelCase_       :					EvalPrediction		) -> Dict:
												lowerCAmelCase__						,       lowerCAmelCase__    						=						align_predictions(p.predictions    ,       p.label_ids		)
												return {
												    "accuracy_score": accuracy_score(UpperCamelCase_    ,       UpperCamelCase_		),
												    "precision": precision_score(UpperCamelCase_    ,       UpperCamelCase_		),
												    "recall": recall_score(UpperCamelCase_    ,       UpperCamelCase_		),
												    "f1": fa_score(UpperCamelCase_    ,       UpperCamelCase_		),
												}

						# Data collator
						lowerCAmelCase__    						=						DataCollatorWithPadding(UpperCamelCase_    ,       pad_to_multiple_of=8		) if training_args.fpaa else None

						# Initialize our Trainer
						lowerCAmelCase__    						=						Trainer(
						    model=UpperCamelCase_    ,       args=UpperCamelCase_    ,       train_dataset=UpperCamelCase_    ,       eval_dataset=UpperCamelCase_    ,       compute_metrics=UpperCamelCase_    ,       data_collator=UpperCamelCase_    ,       )

						# Training
						if training_args.do_train:
												trainer.train(
												    model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path		) else None		)
												trainer.save_model()
												# For convenience, we also re-save the tokenizer to the same directory,
												# so that you can share your model easily on huggingface.co/models =)
												if trainer.is_world_process_zero():
																		tokenizer.save_pretrained(training_args.output_dir		)

    # Evaluation
						lowerCAmelCase__    						=						{}
						if training_args.do_eval:
												logger.info('*** Evaluate ***'		)

												lowerCAmelCase__    						=						trainer.evaluate()

												lowerCAmelCase__    						=						os.path.join(training_args.output_dir    ,       'eval_results.txt'		)
												if trainer.is_world_process_zero():
																		with open(UpperCamelCase_    ,       'w'		) as writer:
																								logger.info('***** Eval results *****'		)
																								for key, value in result.items():
																														logger.info('  %s = %s'    ,       UpperCamelCase_    ,       UpperCamelCase_		)
																														writer.write('%s = %s\n' % (key, value)		)

																		results.update(UpperCamelCase_		)

    # Predict
						if training_args.do_predict:
												lowerCAmelCase__    						=						TokenClassificationDataset(
												    token_classification_task=UpperCamelCase_    ,       data_dir=data_args.data_dir    ,       tokenizer=UpperCamelCase_    ,       labels=UpperCamelCase_    ,       model_type=config.model_type    ,       max_seq_length=data_args.max_seq_length    ,       overwrite_cache=data_args.overwrite_cache    ,       mode=Split.test    ,       )

												lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__    						=						trainer.predict(UpperCamelCase_		)
												lowerCAmelCase__						,       lowerCAmelCase__    						=						align_predictions(UpperCamelCase_    ,       UpperCamelCase_		)

												lowerCAmelCase__    						=						os.path.join(training_args.output_dir    ,       'test_results.txt'		)
												if trainer.is_world_process_zero():
																		with open(UpperCamelCase_    ,       'w'		) as writer:
																								for key, value in metrics.items():
																														logger.info('  %s = %s'    ,       UpperCamelCase_    ,       UpperCamelCase_		)
																														writer.write('%s = %s\n' % (key, value)		)

        # Save predictions
												lowerCAmelCase__    						=						os.path.join(training_args.output_dir    ,       'test_predictions.txt'		)
												if trainer.is_world_process_zero():
																		with open(UpperCamelCase_    ,       'w'		) as writer:
																								with open(os.path.join(data_args.data_dir    ,       'test.txt'		)    ,       'r'		) as f:
																														token_classification_task.write_predictions_to_file(UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_		)

						return results







def 	_UpperCamelCase   (    UpperCamelCase_       :					Dict		)  ->   Union[str, Any]:

						"""simple docstring"""



						main()


if __name__ == "__main__":
					main()





