'''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__     :	str				=			True
  from torch.cuda.amp import autocast


UpperCamelCase__     :	int				=			logging.getLogger(__name__)


@dataclass
class 				_lowercase	:
      '''simple docstring'''


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







def __UpperCamelCase(   _A					:				ModelArguments  ,			_A					:				TrainingArguments    ):
   '''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    )]  ,			)
   UpperCAmelCase__							:				List[Any]   				=  logging.WARNING
   if model_args.verbose_logging:
      UpperCAmelCase__							:				Union[str, Any]   				=  logging.DEBUG
   elif trainer_utils.is_main_process(training_args.local_rank    ):
      UpperCAmelCase__							:				Optional[Any]   				=  logging.INFO
   logger.setLevel(_A    )


@dataclass
class 				_lowercase	:
      '''simple docstring'''


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




@dataclass
class 				_lowercase	:
      '''simple docstring'''


      UpperCAmelCase_      :      WavaVecaForPreTraining
      UpperCAmelCase_      :      WavaVecaFeatureExtractor
      UpperCAmelCase_      :      Union[bool, str]      							=   "longest"
      UpperCAmelCase_      :      Optional[int]      							=   None
      UpperCAmelCase_      :      Optional[int]      							=   None




      def __call__(				self					,lowerCamelCase_							)			->			Dict[str, torch.Tensor]:



         '''simple docstring'''



         UpperCAmelCase__							:				Any   				=  self.feature_extractor.pad(
             lowerCamelCase_					,max_length=self.max_length					,padding=self.padding					,pad_to_multiple_of=self.pad_to_multiple_of					,return_tensors='''pt'''					,)
         UpperCAmelCase__							:				Tuple   				=  self.model._get_feat_extract_output_lengths(batch['''input_values'''].shape[-1]							)

         UpperCAmelCase__							:				List[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
            UpperCAmelCase__							:				List[str]   				=  self.model._get_feat_extract_output_lengths(batch['''attention_mask'''].sum(-1							)							).to(
                torch.long							)

            UpperCAmelCase__							:				Optional[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
            UpperCAmelCase__							:				str   				=  1
            UpperCAmelCase__							:				str   				=  attention_mask.flip([-1]							).cumsum(-1							).flip([-1]							).bool()

         # sample randomly masked indices
         UpperCAmelCase__							:				Any   				=  _compute_mask_indices(
             (batch_size, mask_indices_seq_length)					,self.model.config.mask_time_prob					,self.model.config.mask_time_length					,attention_mask=lowerCamelCase_					,min_masks=2					,)

         return batch




class 				_lowercase	(      lowerCAmelCase   ):
      '''simple docstring'''


      def __init__(				self					,*lowerCamelCase_					,lowerCamelCase_=1					,lowerCamelCase_=0					,lowerCamelCase_=1.0					,**lowerCamelCase_							)			->			List[str]:



         '''simple docstring'''



         super().__init__(*lowerCamelCase_					,**lowerCamelCase_							)
         UpperCAmelCase__							:				Any   				=  0
         UpperCAmelCase__							:				List[str]   				=  max_gumbel_temp
         UpperCAmelCase__							:				List[Any]   				=  min_gumbel_temp
         UpperCAmelCase__							:				Optional[Any]   				=  gumbel_temp_decay




      def 		lowerCAmelCase__					(				self					,lowerCamelCase_					,lowerCamelCase_							)			->			torch.Tensor:



         '''simple docstring'''



         model.train()
         UpperCAmelCase__							:				Union[str, Any]   				=  self._prepare_inputs(lowerCamelCase_							)

         if self.use_amp:
            with autocast():
               UpperCAmelCase__							:				Any   				=  self.compute_loss(lowerCamelCase_					,lowerCamelCase_							)
         else:
            UpperCAmelCase__							:				Optional[int]   				=  self.compute_loss(lowerCamelCase_					,lowerCamelCase_							)

         if self.args.n_gpu > 1 or self.deepspeed:
            if model.module.config.ctc_loss_reduction == "mean":
               UpperCAmelCase__							:				List[Any]   				=  loss.mean()
            elif model.module.config.ctc_loss_reduction == "sum":
               UpperCAmelCase__							:				List[Any]   				=  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:
            UpperCAmelCase__							:				List[str]   				=  loss / self.args.gradient_accumulation_steps

         if self.use_amp:
            self.scaler.scale(lowerCamelCase_							).backward()
         elif self.use_apex:
            with amp.scale_loss(lowerCamelCase_					,self.optimizer							) as scaled_loss:
               scaled_loss.backward()
         elif self.deepspeed:
            self.deepspeed.backward(lowerCamelCase_							)
         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 __UpperCamelCase(   ):
   '''simple docstring'''
   # See all possible arguments in src/transformers/training_args.py
   # or by passing the --help flag to this script.
   # We now keep distinct sets of args, for a cleaner separation of concerns.

   UpperCAmelCase__							:				Dict   				=  HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)    )

   UpperCAmelCase__				,      UpperCAmelCase__				,      UpperCAmelCase__							:				Optional[int]   				=  parser.parse_args_into_dataclasses()
   configure_logger(_A  ,			_A    )

   # Downloading and loading a dataset from the hub.
   UpperCAmelCase__							:				Optional[int]   				=  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"
      UpperCAmelCase__							:				Any   				=  DatasetDict()
      UpperCAmelCase__							:				List[str]   				=  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  ,			)
      UpperCAmelCase__							:				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"
      UpperCAmelCase__							:				Any   				=  DatasetDict()
      UpperCAmelCase__							:				Optional[Any]   				=  load_dataset(
          data_args.dataset_name  ,			data_args.dataset_config_name  ,			split='''validation'''  ,			cache_dir=model_args.cache_dir  ,			)
      UpperCAmelCase__							:				Dict   				=  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
   UpperCAmelCase__							:				Any   				=  WavaVecaFeatureExtractor.from_pretrained(
       model_args.model_name_or_path  ,			cache_dir=model_args.cache_dir  ,			do_normalize=_A    )

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

   # load audio files into numpy arrays
   UpperCAmelCase__							:				Union[str, Any]   				=  datasets.map(
       _A  ,			num_proc=data_args.preprocessing_num_workers  ,			remove_columns=datasets['''train'''].column_names    )

   # filter audio files that are too long
   UpperCAmelCase__							:				List[str]   				=  vectorized_datasets.filter(
       lambda _A    : len(data['''speech''']    ) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate    )    )

   def normalize(_A					:				Dict    ):
      return feature_extractor(batch['''speech''']  ,			sampling_rate=feature_extractor.sampling_rate    )

   # normalize and transform to `BatchFeatures`
   UpperCAmelCase__							:				Optional[Any]   				=  vectorized_datasets.map(
       _A  ,			batched=_A  ,			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
   UpperCAmelCase__							:				Tuple   				=  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\''''    )

   UpperCAmelCase__							:				int   				=  WavaVecaForPreTraining(_A    )

   UpperCAmelCase__							:				Tuple   				=  DataCollatorForWavaVecaPretraining(model=_A  ,			feature_extractor=_A    )

   UpperCAmelCase__							:				Optional[int]   				=  WavaVecaPreTrainer(
       model=_A  ,			data_collator=_A  ,			args=_A  ,			train_dataset=vectorized_datasets['''train''']  ,			eval_dataset=vectorized_datasets['''validation''']  ,			tokenizer=_A  ,			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


UpperCamelCase__     :	List[Any]				=			logging.get_logger(__name__)

UpperCamelCase__     :	Optional[Any]				=			{
    '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 				_lowercase	(      lowerCAmelCase   ):
      '''simple docstring'''


      UpperCAmelCase_      :      Any      							=   '''time_series_transformer'''
      UpperCAmelCase_      :      List[Any]      							=   {
          '''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_ = 32					,lowerCamelCase_ = 32					,lowerCamelCase_ = 2					,lowerCamelCase_ = 2					,lowerCamelCase_ = 2					,lowerCamelCase_ = 2					,lowerCamelCase_ = True					,lowerCamelCase_ = "gelu"					,lowerCamelCase_ = 64					,lowerCamelCase_ = 0.1					,lowerCamelCase_ = 0.1					,lowerCamelCase_ = 0.1					,lowerCamelCase_ = 0.1					,lowerCamelCase_ = 0.1					,lowerCamelCase_ = 100					,lowerCamelCase_ = 0.02					,lowerCamelCase_=True					,**lowerCamelCase_					,)			->			Tuple:



         '''simple docstring'''



         UpperCAmelCase__							:				Dict   				=  prediction_length
         UpperCAmelCase__							:				Optional[Any]   				=  context_length or prediction_length
         UpperCAmelCase__							:				Union[str, Any]   				=  distribution_output
         UpperCAmelCase__							:				List[Any]   				=  loss
         UpperCAmelCase__							:				Optional[int]   				=  input_size
         UpperCAmelCase__							:				Tuple   				=  num_time_features
         UpperCAmelCase__							:				str   				=  lags_sequence
         UpperCAmelCase__							:				Union[str, Any]   				=  scaling
         UpperCAmelCase__							:				Optional[int]   				=  num_dynamic_real_features
         UpperCAmelCase__							:				Optional[Any]   				=  num_static_real_features
         UpperCAmelCase__							:				str   				=  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`'''							)
            UpperCAmelCase__							:				Union[str, Any]   				=  cardinality
         else:
            UpperCAmelCase__							:				List[str]   				=  [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`'''							)
            UpperCAmelCase__							:				int   				=  embedding_dimension
         else:
            UpperCAmelCase__							:				Union[str, Any]   				=  [min(50					,(cat + 1) // 2							) for cat in self.cardinality]
         UpperCAmelCase__							:				Dict   				=  num_parallel_samples

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

         UpperCAmelCase__							:				Optional[Any]   				=  dropout
         UpperCAmelCase__							:				Any   				=  attention_dropout
         UpperCAmelCase__							:				Optional[int]   				=  activation_dropout
         UpperCAmelCase__							:				List[Any]   				=  encoder_layerdrop
         UpperCAmelCase__							:				Optional[int]   				=  decoder_layerdrop

         UpperCAmelCase__							:				List[Any]   				=  activation_function
         UpperCAmelCase__							:				List[Any]   				=  init_std

         UpperCAmelCase__							:				Optional[Any]   				=  use_cache

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




      @property
      def 		lowerCAmelCase__					(				self							)			->			int:



         '''simple docstring'''



         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
         )




