





'''simple docstring'''






import logging
import sys
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

import librosa
import torch
from datasets import DatasetDict, load_dataset
from packaging import version
from torch import nn

from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaForPreTraining,
    is_apex_available,
    trainer_utils,
)
from transformers.models.wavaveca.modeling_wavaveca import _compute_mask_indices


if is_apex_available():
	from apex import amp

if version.parse(version.parse(torch.__version__).base_version) >= version.parse("1.6"):
	__UpperCamelCase					    =	True
	from torch.cuda.amp import autocast


__UpperCamelCase					    =	logging.getLogger(__name__)




@dataclass
class  _A :
		lowercase__:   str				     =   field(
		    metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models'''}		)
		lowercase__:   Optional[str]				     =   field(
		    default=__lowercase						,					metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co'''}						,					)
		lowercase__:   Optional[bool]				     =   field(
		    default=__lowercase						,					metadata={'''help''': '''Whether to freeze the feature extractor layers of the model.'''}		)
		lowercase__:   Optional[bool]				     =   field(
		    default=__lowercase						,					metadata={'''help''': '''Whether to log verbose messages or not.'''}						,					)
		lowercase__:   Optional[float]				     =   field(
		    default=2.0						,					metadata={'''help''': '''Maximum temperature for gumbel softmax.'''}		)
		lowercase__:   Optional[float]				     =   field(
		    default=0.5						,					metadata={'''help''': '''Minimum temperature for gumbel softmax.'''}		)
		lowercase__:   Optional[float]				     =   field(
		    default=0.9_9_9_9_9_5						,					metadata={'''help''': '''Decay of gumbel temperature during training.'''}		)







def 					_a						(       _lowerCamelCase    ,	_lowerCamelCase ) ->				Optional[int]:



			"""simple docstring"""


			logging.basicConfig(
			    format="""%(asctime)s - %(levelname)s - %(name)s - %(message)s"""    ,	datefmt="""%m/%d/%Y %H:%M:%S"""    ,	handlers=[logging.StreamHandler(sys.stdout )]    ,	)
			__snake_case					:      int							      =     logging.WARNING
			if model_args.verbose_logging:
						__snake_case					:      Any							      =     logging.DEBUG
			elif trainer_utils.is_main_process(training_args.local_rank ):
						__snake_case					:      List[str]							      =     logging.INFO
			logger.setLevel(_lowerCamelCase )




@dataclass
class  _A :
		lowercase__:   str				     =   field(
		    default=__lowercase						,					metadata={'''help''': '''The name of the dataset to use (via the datasets library).'''}		)
		lowercase__:   Optional[str]				     =   field(
		    default=__lowercase						,					metadata={'''help''': '''The configuration name of the dataset to use (via the datasets library).'''}		)
		lowercase__:   Optional[str]				     =   field(
		    default='''train'''						,					metadata={
		        '''help''': '''The name of the training data set split to use (via the datasets library). Defaults to \'train\''''
		    }						,					)
		lowercase__:   Optional[str]				     =   field(
		    default='''validation'''						,					metadata={
		        '''help''': (
		            '''The name of the validation data set split to use (via the datasets library). Defaults to \'validation\''''
		        )
		    }						,					)
		lowercase__:   Optional[str]				     =   field(
		    default='''file'''						,					metadata={'''help''': '''Column in the dataset that contains speech file path. Defaults to \'file\''''}						,					)
		lowercase__:   bool				     =   field(
		    default=__lowercase						,					metadata={'''help''': '''Overwrite the cached preprocessed datasets or not.'''}		)
		lowercase__:   Optional[int]				     =   field(
		    default=1						,					metadata={
		        '''help''': '''The percentage of the train set used as validation set in case there\'s no validation split'''
		    }						,					)
		lowercase__:   Optional[int]				     =   field(
		    default=__lowercase						,					metadata={'''help''': '''The number of processes to use for the preprocessing.'''}						,					)
		lowercase__:   Optional[float]				     =   field(
		    default=2_0.0						,					metadata={'''help''': '''Filter audio files that are longer than `max_duration_in_seconds` seconds'''}		)





@dataclass
class  _A :
		lowercase__:   WavaVecaForPreTraining
		lowercase__:   WavaVecaFeatureExtractor
		lowercase__:   Union[bool, str]				     =   "longest"
		lowercase__:   Optional[int]				     =   None
		lowercase__:   Optional[int]				     =   None

		def __call__(   self					:   Any  ,      __magic_name__					:   List[Dict[str, Union[List[int], torch.Tensor]]]			)					->    Dict[str, torch.Tensor]:


					"""simple docstring"""

					__snake_case					:      Any							      =     self.feature_extractor.pad(
					    __magic_name__  ,      max_length=self.max_length  ,      padding=self.padding  ,      pad_to_multiple_of=self.pad_to_multiple_of  ,      return_tensors="""pt"""  ,      )
					__snake_case					:      int							      =     self.model._get_feat_extract_output_lengths(batch["""input_values"""].shape[-1]			)

					__snake_case					:      Optional[Any]							      =     batch["""input_values"""].shape[0]

					# make sure that no loss is computed on padded inputs
					if batch["attention_mask"] is not None:
								# compute real output lengths according to convolution formula
								__snake_case					:      Optional[int]							      =     self.model._get_feat_extract_output_lengths(batch["""attention_mask"""].sum(-1			)			).to(
								    torch.long			)

								__snake_case					:      Any							      =     torch.zeros(
								    (batch_size, mask_indices_seq_length)  ,      dtype=torch.long  ,      device=batch["""input_values"""].device			)

								# these two operations makes sure that all values
								# before the output lengths indices are attended to
								__snake_case					:      Dict							      =     1
								__snake_case					:      Optional[Any]							      =     attention_mask.flip([-1]			).cumsum(-1			).flip([-1]			).bool()

					# sample randomly masked indices
					__snake_case					:      int							      =     _compute_mask_indices(
					    (batch_size, mask_indices_seq_length)  ,      self.model.config.mask_time_prob  ,      self.model.config.mask_time_length  ,      attention_mask=__magic_name__  ,      min_masks=2  ,      )

					return batch





class  _A (					__lowercase		):

		def __init__(   self					:   str  ,      *__magic_name__					:   Any  ,      __magic_name__					:   Optional[Any]=1  ,      __magic_name__					:   int=0  ,      __magic_name__					:   Tuple=1.0  ,      **__magic_name__					:   Any			)					->    Dict:


					"""simple docstring"""

					super().__init__(*__magic_name__  ,      **__magic_name__			)
					__snake_case					:      List[str]							      =     0
					__snake_case					:      Optional[Any]							      =     max_gumbel_temp
					__snake_case					:      int							      =     min_gumbel_temp
					__snake_case					:      Union[str, Any]							      =     gumbel_temp_decay

		def 	lowercase__     (   self					:   Optional[Any]  ,      __magic_name__					:   nn.Module  ,      __magic_name__					:   Dict[str, Union[torch.Tensor, Any]]			)					->    torch.Tensor:


					"""simple docstring"""

					model.train()
					__snake_case					:      int							      =     self._prepare_inputs(__magic_name__			)

					if self.use_amp:
								with autocast():
											__snake_case					:      str							      =     self.compute_loss(__magic_name__  ,      __magic_name__			)
					else:
								__snake_case					:      Tuple							      =     self.compute_loss(__magic_name__  ,      __magic_name__			)

					if self.args.n_gpu > 1 or self.deepspeed:
								if model.module.config.ctc_loss_reduction == "mean":
											__snake_case					:      Union[str, Any]							      =     loss.mean()
								elif model.module.config.ctc_loss_reduction == "sum":
											__snake_case					:      str							      =     loss.sum() / (inputs["""mask_time_indices"""]).sum()
								else:
											raise ValueError(f'''{model.config.ctc_loss_reduction} is not valid. Choose one of [\'mean\', \'sum\']'''			)

					if self.args.gradient_accumulation_steps > 1:
								__snake_case					:      List[Any]							      =     loss / self.args.gradient_accumulation_steps

					if self.use_amp:
								self.scaler.scale(__magic_name__			).backward()
					elif self.use_apex:
								with amp.scale_loss(__magic_name__  ,      self.optimizer			) as scaled_loss:
											scaled_loss.backward()
					elif self.deepspeed:
								self.deepspeed.backward(__magic_name__			)
					else:
								loss.backward()

					self.num_update_step += 1
					# make sure gumbel softmax temperature is decayed
					if self.args.n_gpu > 1 or self.deepspeed:
								model.module.set_gumbel_temperature(
								    max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step  ,      self.min_gumbel_temp			)			)
					else:
								model.set_gumbel_temperature(
								    max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step  ,      self.min_gumbel_temp			)			)

					return loss.detach()







def 					_a						(       ) ->				Any:



			"""simple docstring"""


			__snake_case					:      Union[str, Any]							      =     HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments) )

			__snake_case      ,     __snake_case      ,     __snake_case					:      List[Any]							      =     parser.parse_args_into_dataclasses()
			configure_logger(_lowerCamelCase    ,	_lowerCamelCase )

			# Downloading and loading a dataset from the hub.
			__snake_case					:      Tuple							      =     load_dataset(data_args.dataset_name    ,	data_args.dataset_config_name    ,	cache_dir=model_args.cache_dir )

			if "validation" not in datasets.keys():
						# make sure only "validation" and "train" keys remain"
						__snake_case					:      int							      =     DatasetDict()
						__snake_case					:      Union[str, Any]							      =     load_dataset(
						    data_args.dataset_name    ,	data_args.dataset_config_name    ,	split=F'''{data_args.train_split_name}[:{data_args.validation_split_percentage}%]'''    ,	cache_dir=model_args.cache_dir    ,	)
						__snake_case					:      Optional[int]							      =     load_dataset(
						    data_args.dataset_name    ,	data_args.dataset_config_name    ,	split=F'''{data_args.train_split_name}[{data_args.validation_split_percentage}%:]'''    ,	cache_dir=model_args.cache_dir    ,	)
			else:
						# make sure only "validation" and "train" keys remain"
						__snake_case					:      Optional[Any]							      =     DatasetDict()
						__snake_case					:      Optional[Any]							      =     load_dataset(
						    data_args.dataset_name    ,	data_args.dataset_config_name    ,	split="""validation"""    ,	cache_dir=model_args.cache_dir    ,	)
						__snake_case					:      str							      =     load_dataset(
						    data_args.dataset_name    ,	data_args.dataset_config_name    ,	split=F'''{data_args.train_split_name}'''    ,	cache_dir=model_args.cache_dir    ,	)

			# only normalized-inputs-training is supported
			__snake_case					:      Union[str, Any]							      =     WavaVecaFeatureExtractor.from_pretrained(
			    model_args.model_name_or_path    ,	cache_dir=model_args.cache_dir    ,	do_normalize=_lowerCamelCase )

			def prepare_dataset(_lowerCamelCase ):
						# check that all files have the correct sampling rate
						__snake_case      ,     __snake_case					:      int							      =     librosa.load(batch[data_args.speech_file_column]    ,	sr=feature_extractor.sampling_rate )
						return batch

			# load audio files into numpy arrays
			__snake_case					:      List[str]							      =     datasets.map(
			    _lowerCamelCase    ,	num_proc=data_args.preprocessing_num_workers    ,	remove_columns=datasets["""train"""].column_names )

			# filter audio files that are too long
			__snake_case					:      Optional[Any]							      =     vectorized_datasets.filter(
			    lambda _lowerCamelCase : len(data["""speech"""] ) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate ) )

			def normalize(_lowerCamelCase ):
						return feature_extractor(batch["""speech"""]    ,	sampling_rate=feature_extractor.sampling_rate )

			# normalize and transform to `BatchFeatures`
			__snake_case					:      str							      =     vectorized_datasets.map(
			    _lowerCamelCase    ,	batched=_lowerCamelCase    ,	num_proc=data_args.preprocessing_num_workers    ,	load_from_cache_file=not data_args.overwrite_cache    ,	remove_columns=vectorized_datasets["""train"""].column_names    ,	)

			# pretraining is only supported for "newer" stable layer norm architecture
			# apply_spec_augment has to be True, mask_feature_prob has to be 0.0
			__snake_case					:      str							      =     WavaVecaConfig.from_pretrained(
			    model_args.model_name_or_path    ,	cache_dir=model_args.cache_dir    ,	gradient_checkpointing=training_args.gradient_checkpointing    ,	)

			if not config.do_stable_layer_norm or config.feat_extract_norm != "layer":
						raise ValueError(
						    """PreTraining is only supported for ``config.do_stable_layer_norm=True`` and"""
						    """ ``config.feat_extract_norm='layer'""" )

			__snake_case					:      Tuple							      =     WavaVecaForPreTraining(_lowerCamelCase )

			__snake_case					:      int							      =     DataCollatorForWavaVecaPretraining(model=_lowerCamelCase    ,	feature_extractor=_lowerCamelCase )

			__snake_case					:      List[str]							      =     WavaVecaPreTrainer(
			    model=_lowerCamelCase    ,	data_collator=_lowerCamelCase    ,	args=_lowerCamelCase    ,	train_dataset=vectorized_datasets["""train"""]    ,	eval_dataset=vectorized_datasets["""validation"""]    ,	tokenizer=_lowerCamelCase    ,	max_gumbel_temp=model_args.max_gumbel_temperature    ,	min_gumbel_temp=model_args.min_gumbel_temperature    ,	gumbel_temp_decay=model_args.gumbel_temperature_decay    ,	)
			trainer.train()


if __name__ == "__main__":
	main()










"""simple docstring"""
from typing import List, Optional, Union

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


lowerCAmelCase__    						=			logging.get_logger(__name__)

lowerCAmelCase__    						=			{
    'huggingface/time-series-transformer-tourism-monthly': (
        'https://huggingface.co/huggingface/time-series-transformer-tourism-monthly/resolve/main/config.json'
    ),
    # See all TimeSeriesTransformer models at https://huggingface.co/models?filter=time_series_transformer
}


class _lowerCAmelCase						(    __UpperCAmelCase       ):
						SCREAMING_SNAKE_CASE_:					Any 					=				'time_series_transformer'
						SCREAMING_SNAKE_CASE_:					List[str] 					=				{
						    'hidden_size': 'd_model',
						    'num_attention_heads': 'encoder_attention_heads',
						    'num_hidden_layers': 'encoder_layers',
						}







						def __init__(    self   ,				lowerCAmelCase_ = None   ,				lowerCAmelCase_ = None   ,				lowerCAmelCase_ = "student_t"   ,				lowerCAmelCase_ = "nll"   ,				lowerCAmelCase_ = 1   ,				lowerCAmelCase_ = [1, 2, 3, 4, 5, 6, 7]   ,				lowerCAmelCase_ = "mean"   ,				lowerCAmelCase_ = 0   ,				lowerCAmelCase_ = 0   ,				lowerCAmelCase_ = 0   ,				lowerCAmelCase_ = 0   ,				lowerCAmelCase_ = None   ,				lowerCAmelCase_ = None   ,				lowerCAmelCase_ = 3_2   ,				lowerCAmelCase_ = 3_2   ,				lowerCAmelCase_ = 2   ,				lowerCAmelCase_ = 2   ,				lowerCAmelCase_ = 2   ,				lowerCAmelCase_ = 2   ,				lowerCAmelCase_ = True   ,				lowerCAmelCase_ = "gelu"   ,				lowerCAmelCase_ = 6_4   ,				lowerCAmelCase_ = 0.1   ,				lowerCAmelCase_ = 0.1   ,				lowerCAmelCase_ = 0.1   ,				lowerCAmelCase_ = 0.1   ,				lowerCAmelCase_ = 0.1   ,				lowerCAmelCase_ = 1_0_0   ,				lowerCAmelCase_ = 0.02   ,				lowerCAmelCase_=True   ,				**lowerCAmelCase_   ,				)		->							str:
										# time series specific configuration
										_SCREAMING_SNAKE_CASE   :		List[str]       		=    prediction_length
										_SCREAMING_SNAKE_CASE   :		Tuple       		=    context_length or prediction_length
										_SCREAMING_SNAKE_CASE   :		Any       		=    distribution_output
										_SCREAMING_SNAKE_CASE   :		List[str]       		=    loss
										_SCREAMING_SNAKE_CASE   :		Any       		=    input_size
										_SCREAMING_SNAKE_CASE   :		Any       		=    num_time_features
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    lags_sequence
										_SCREAMING_SNAKE_CASE   :		Dict       		=    scaling
										_SCREAMING_SNAKE_CASE   :		Any       		=    num_dynamic_real_features
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    num_static_real_features
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    num_static_categorical_features
										if cardinality and num_static_categorical_features > 0:
														if len(lowerCAmelCase_							) != num_static_categorical_features:
																		raise ValueError(
																		    'The cardinality should be a list of the same length as `num_static_categorical_features`'							)
														_SCREAMING_SNAKE_CASE   :		List[str]       		=    cardinality
										else:
														_SCREAMING_SNAKE_CASE   :		Dict       		=    [0]
										if embedding_dimension and num_static_categorical_features > 0:
														if len(lowerCAmelCase_							) != num_static_categorical_features:
																		raise ValueError(
																		    'The embedding dimension should be a list of the same length as `num_static_categorical_features`'							)
														_SCREAMING_SNAKE_CASE   :		int       		=    embedding_dimension
										else:
														_SCREAMING_SNAKE_CASE   :		List[str]       		=    [min(5_0   ,				(cat + 1) // 2							) for cat in self.cardinality]
										_SCREAMING_SNAKE_CASE   :		str       		=    num_parallel_samples

										# Transformer architecture configuration
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    input_size * len(lowerCAmelCase_							) + self._number_of_features
										_SCREAMING_SNAKE_CASE   :		int       		=    d_model
										_SCREAMING_SNAKE_CASE   :		str       		=    encoder_attention_heads
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    decoder_attention_heads
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    encoder_ffn_dim
										_SCREAMING_SNAKE_CASE   :		Dict       		=    decoder_ffn_dim
										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    encoder_layers
										_SCREAMING_SNAKE_CASE   :		Tuple       		=    decoder_layers

										_SCREAMING_SNAKE_CASE   :		Dict       		=    dropout
										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    attention_dropout
										_SCREAMING_SNAKE_CASE   :		Tuple       		=    activation_dropout
										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    encoder_layerdrop
										_SCREAMING_SNAKE_CASE   :		List[str]       		=    decoder_layerdrop

										_SCREAMING_SNAKE_CASE   :		Dict       		=    activation_function
										_SCREAMING_SNAKE_CASE   :		str       		=    init_std

										_SCREAMING_SNAKE_CASE   :		Tuple       		=    use_cache

										super().__init__(is_encoder_decoder=lowerCAmelCase_   ,				**lowerCAmelCase_							)







						@property
						def 			A      (    self							)		->							int:
										return (
										    sum(self.embedding_dimension							)
										    + self.num_dynamic_real_features
										    + self.num_time_features
										    + self.num_static_real_features
										    + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
										)

