



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"):
       a__   :					Optional[Any] 						=      True
       from torch.cuda.amp import autocast


a__   :					List[str] 						=      logging.getLogger(__name__)



@dataclass
class    UpperCAmelCase__:

   '''simple docstring'''




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

def 	_lowerCAmelCase (    A__    ,	A__						):
      logging.basicConfig(
          format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'    ,	datefmt='%m/%d/%Y %H:%M:%S'    ,	handlers=[logging.StreamHandler(sys.stdout						)]    ,	)
      lowercase__  =	logging.WARNING
      if model_args.verbose_logging:
            lowercase__  =	logging.DEBUG
      elif trainer_utils.is_main_process(training_args.local_rank						):
            lowercase__  =	logging.INFO
      logger.setLevel(A__						)



@dataclass
class    UpperCAmelCase__:

   '''simple docstring'''




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



@dataclass
class    UpperCAmelCase__:

   '''simple docstring'''




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


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




         """simple docstring"""

         lowercase__  =	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'      ,					)
         lowercase__  =	self.model._get_feat_extract_output_lengths(batch['input_values'].shape[-1])

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

               lowercase__  =	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
               lowercase__  =	1
               lowercase__  =	attention_mask.flip([-1]).cumsum(-1).flip([-1]).bool()

         # sample randomly masked indices
         lowercase__  =	_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    UpperCAmelCase__(				lowerCamelCase						):

   '''simple docstring'''




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




         """simple docstring"""

         super().__init__(*lowerCAmelCase      ,					**lowerCAmelCase)
         lowercase__  =	0
         lowercase__  =	max_gumbel_temp
         lowercase__  =	min_gumbel_temp
         lowercase__  =	gumbel_temp_decay



   def       UpperCAmelCase		(    self	:      List[str]      ,					lowerCAmelCase	:      nn.Module      ,					lowerCAmelCase	:      Dict[str, Union[torch.Tensor, Any]])						->					torch.Tensor:




         """simple docstring"""

         model.train()
         lowercase__  =	self._prepare_inputs(lowerCAmelCase)

         if self.use_amp:
               with autocast():
                     lowercase__  =	self.compute_loss(lowerCAmelCase      ,					lowerCAmelCase)
         else:
               lowercase__  =	self.compute_loss(lowerCAmelCase      ,					lowerCAmelCase)

         if self.args.n_gpu > 1 or self.deepspeed:
               if model.module.config.ctc_loss_reduction == "mean":
                     lowercase__  =	loss.mean()
               elif model.module.config.ctc_loss_reduction == "sum":
                     lowercase__  =	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:
               lowercase__  =	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 	_lowerCAmelCase (    ):
      # 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.

      lowercase__  =	HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)						)

      lowercase__,						lowercase__,						lowercase__  =	parser.parse_args_into_dataclasses()
      configure_logger(A__    ,	A__						)

      # Downloading and loading a dataset from the hub.
      lowercase__  =	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"
            lowercase__  =	DatasetDict()
            lowercase__  =	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    ,	)
            lowercase__  =	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"
            lowercase__  =	DatasetDict()
            lowercase__  =	load_dataset(
                data_args.dataset_name    ,	data_args.dataset_config_name    ,	split='validation'    ,	cache_dir=model_args.cache_dir    ,	)
            lowercase__  =	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
      lowercase__  =	WavaVecaFeatureExtractor.from_pretrained(
          model_args.model_name_or_path    ,	cache_dir=model_args.cache_dir    ,	do_normalize=A__						)

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

      # load audio files into numpy arrays
      lowercase__  =	datasets.map(
          A__    ,	num_proc=data_args.preprocessing_num_workers    ,	remove_columns=datasets['train'].column_names						)

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

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

      # normalize and transform to `BatchFeatures`
      lowercase__  =	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
      lowercase__  =	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\''						)

      lowercase__  =	WavaVecaForPreTraining(A__						)

      lowercase__  =	DataCollatorForWavaVecaPretraining(model=A__    ,	feature_extractor=A__						)

      lowercase__  =	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()





a__   :					Tuple 						=      "Tobias Carryer"

from time import time



class    UpperCAmelCase__:

   '''simple docstring'''




   def __init__(    self	:      str      ,					lowerCAmelCase	:      List[str]      ,					lowerCAmelCase	:      Any      ,					lowerCAmelCase	:      str      ,					lowerCAmelCase	:      str=int(time()))						->					List[Any]:  # noqa: B008




         """simple docstring"""

         lowercase__  =	multiplier
         lowercase__  =	increment
         lowercase__  =	modulo
         lowercase__  =	seed



   def       UpperCAmelCase		(    self	:      int)						->					Tuple:




         """simple docstring"""

         lowercase__  =	(self.multiplier * self.seed + self.increment) % self.modulo
         return self.seed


if __name__ == "__main__":
       # Show the LCG in action.
       a__   :					str 						=      LinearCongruentialGenerator(1_66_45_25, 10_13_90_42_23, 2 << 31)
       while True:
              print(lcg.next_number())
