



"""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"""):
  lowerCamelCase  			=		True
  from torch.cuda.amp import autocast


lowerCamelCase  			=		logging.getLogger(__name__)
@dataclass
class  lowercase__  :







       '''simple docstring'''


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



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







       '''simple docstring'''


       UpperCamelCase		=	field(
           default=SCREAMING_SNAKE_CASE							,   metadata={'''help''': '''The name of the dataset to use (via the datasets library).'''} )
       UpperCamelCase		=	field(
           default=SCREAMING_SNAKE_CASE							,   metadata={'''help''': '''The configuration name of the dataset to use (via the datasets library).'''} )
       UpperCamelCase		=	field(
           default='''train'''							,   metadata={
               '''help''': '''The name of the training data set split to use (via the datasets library). Defaults to \'train\''''
           }							,   )
       UpperCamelCase		=	field(
           default='''validation'''							,   metadata={
               '''help''': (
                   '''The name of the validation data set split to use (via the datasets library). Defaults to \'validation\''''
               )
           }							,   )
       UpperCamelCase		=	field(
           default='''file'''							,   metadata={'''help''': '''Column in the dataset that contains speech file path. Defaults to \'file\''''}							,   )
       UpperCamelCase		=	field(
           default=SCREAMING_SNAKE_CASE							,   metadata={'''help''': '''Overwrite the cached preprocessed datasets or not.'''} )
       UpperCamelCase		=	field(
           default=1							,   metadata={
               '''help''': '''The percentage of the train set used as validation set in case there\'s no validation split'''
           }							,   )
       UpperCamelCase		=	field(
           default=SCREAMING_SNAKE_CASE							,   metadata={'''help''': '''The number of processes to use for the preprocessing.'''}							,   )
       UpperCamelCase		=	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		=	42
       UpperCamelCase		=	42
       UpperCamelCase		=	"longest"
       UpperCamelCase		=	None
       UpperCamelCase		=	None



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


             '''simple docstring'''




             UpperCAmelCase_     					=				self.feature_extractor.pad(
                 _UpperCAmelCase						,							max_length=self.max_length						,							padding=self.padding						,							pad_to_multiple_of=self.pad_to_multiple_of						,							return_tensors="pt"						,							)
             UpperCAmelCase_     					=				self.model._get_feat_extract_output_lengths(batch["input_values"].shape[-1]       )

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

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

             # sample randomly masked indices
             UpperCAmelCase_     					=				_compute_mask_indices(
                 (batch_size, mask_indices_seq_length)						,							self.model.config.mask_time_prob						,							self.model.config.mask_time_length						,							attention_mask=_UpperCAmelCase						,							min_masks=2						,							)

             return batch



class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


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


             '''simple docstring'''




             super().__init__(*_UpperCAmelCase						,							**_UpperCAmelCase       )
             UpperCAmelCase_     					=				0
             UpperCAmelCase_     					=				max_gumbel_temp
             UpperCAmelCase_     					=				min_gumbel_temp
             UpperCAmelCase_     					=				gumbel_temp_decay







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


             '''simple docstring'''




             model.train()
             UpperCAmelCase_     					=				self._prepare_inputs(_UpperCAmelCase       )

             if self.use_amp:
                   with autocast():
                         UpperCAmelCase_     					=				self.compute_loss(_UpperCAmelCase						,							_UpperCAmelCase       )
             else:
                   UpperCAmelCase_     					=				self.compute_loss(_UpperCAmelCase						,							_UpperCAmelCase       )

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

             if self.use_amp:
                   self.scaler.scale(_UpperCAmelCase       ).backward()
             elif self.use_apex:
                   with amp.scale_loss(_UpperCAmelCase						,							self.optimizer       ) as scaled_loss:
                         scaled_loss.backward()
             elif self.deepspeed:
                   self.deepspeed.backward(_UpperCAmelCase       )
             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__      ( ):
      # 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_     					=				HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments) )

      UpperCAmelCase_					,       UpperCAmelCase_					,       UpperCAmelCase_     					=				parser.parse_args_into_dataclasses()
      configure_logger(lowerCAmelCase__      ,     lowerCAmelCase__ )

      # Downloading and loading a dataset from the hub.
      UpperCAmelCase_     					=				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_     					=				DatasetDict()
            UpperCAmelCase_     					=				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_     					=				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_     					=				DatasetDict()
            UpperCAmelCase_     					=				load_dataset(
                data_args.dataset_name      ,     data_args.dataset_config_name      ,     split="validation"      ,     cache_dir=model_args.cache_dir      ,     )
            UpperCAmelCase_     					=				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_     					=				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
            UpperCAmelCase_					,       UpperCAmelCase_     					=				librosa.load(batch[data_args.speech_file_column]      ,     sr=feature_extractor.sampling_rate )
            return batch

      # load audio files into numpy arrays
      UpperCAmelCase_     					=				datasets.map(
          lowerCAmelCase__      ,     num_proc=data_args.preprocessing_num_workers      ,     remove_columns=datasets["train"].column_names )

      # filter audio files that are too long
      UpperCAmelCase_     					=				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`
      UpperCAmelCase_     					=				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
      UpperCAmelCase_     					=				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_     					=				WavaVecaForPreTraining(lowerCAmelCase__ )

      UpperCAmelCase_     					=				DataCollatorForWavaVecaPretraining(model=lowerCAmelCase__      ,     feature_extractor=lowerCAmelCase__ )

      UpperCAmelCase_     					=				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"""






# This code is adapted from OpenAI's release
# https://github.com/openai/human-eval/blob/master/human_eval/execution.py

import contextlib
import faulthandler
import io
import multiprocessing
import os
import platform
import signal
import tempfile




def       A_  (				__lowercase						,	__lowercase						,	__lowercase						,	__lowercase	):
 UpperCamelCase_     :					Tuple   	=multiprocessing.Manager()
 UpperCamelCase_     :					Dict   	=manager.list()

 UpperCamelCase_     :					Union[str, Any]   	=multiprocessing.Process(target=__lowercase						,	args=(check_program, result, timeout)	)
 p.start()
 p.join(timeout=timeout + 1	)
 if p.is_alive():
  p.kill()

 if not result:
  result.append('timed out'	)

 return {
     "task_id": task_id,
     "passed": result[0] == "passed",
     "result": result[0],
     "completion_id": completion_id,
 }




def       A_  (				__lowercase						,	__lowercase						,	__lowercase	):
 with create_tempdir():
  # These system calls are needed when cleaning up tempdir.
  import os
  import shutil

  UpperCamelCase_     :					Tuple   	=shutil.rmtree
  UpperCamelCase_     :					List[str]   	=os.rmdir
  UpperCamelCase_     :					Any   	=os.chdir

  # Disable functionalities that can make destructive changes to the test.
  reliability_guard()

  # Run program.
  try:
   UpperCamelCase_     :					List[str]   	={}
   with swallow_io():
    with time_limit(__lowercase	):
     exec(__lowercase						,	__lowercase	)
   result.append('passed'	)
  except TimeoutException:
   result.append('timed out'	)
  except BaseException as e:
   result.append(F'''failed: {e}'''	)

  # Needed for cleaning up.
  UpperCamelCase_     :					Optional[Any]   	=rmtree
  UpperCamelCase_     :					str   	=rmdir
  UpperCamelCase_     :					Dict   	=chdir




@contextlib.contextmanager
def       A_  (				__lowercase	):
 def signal_handler(__lowercase						,	__lowercase	):
  raise TimeoutException('Timed out!'	)

 signal.setitimer(signal.ITIMER_REAL						,	__lowercase	)
 signal.signal(signal.SIGALRM						,	__lowercase	)
 try:
  yield
 finally:
  signal.setitimer(signal.ITIMER_REAL						,	0	)




@contextlib.contextmanager
def       A_  (				):
 UpperCamelCase_     :					str   	=WriteOnlyStringIO()
 with contextlib.redirect_stdout(__lowercase	):
  with contextlib.redirect_stderr(__lowercase	):
   with redirect_stdin(__lowercase	):
    yield




@contextlib.contextmanager
def       A_  (				):
 with tempfile.TemporaryDirectory() as dirname:
  with chdir(__lowercase	):
   yield dirname





class 				a__					(						A__       ):
     pass





class 				a__					(						io.StringIO       ):



     def  lowerCamelCase_						(							self :Any      ,  *_lowerCamelCase :Any      ,  **_lowerCamelCase :Tuple   ):







      '''simple docstring'''



      raise OSError



     def  lowerCamelCase_						(							self :Union[str, Any]      ,  *_lowerCamelCase :Any      ,  **_lowerCamelCase :Tuple   ):







      '''simple docstring'''



      raise OSError



     def  lowerCamelCase_						(							self :Tuple      ,  *_lowerCamelCase :str      ,  **_lowerCamelCase :List[str]   ):







      '''simple docstring'''



      raise OSError




     def  lowerCamelCase_						(							self :List[str]      ,  *_lowerCamelCase :Optional[int]      ,  **_lowerCamelCase :str   ):







      '''simple docstring'''



      return False





class 				a__					(						contextlib._RedirectStream       ):  # type: ignore
     UpperCAmelCase__         =      '''stdin'''




@contextlib.contextmanager
def       A_  (				__lowercase	):
 if root == ".":
  yield
  return
 UpperCamelCase_     :					List[str]   	=os.getcwd()
 os.chdir(__lowercase	)
 try:
  yield
 except BaseException as exc:
  raise exc
 finally:
  os.chdir(__lowercase	)





def       A_  (				__lowercase=None	):

 if maximum_memory_bytes is not None:
  import resource

  resource.setrlimit(resource.RLIMIT_AS						,	(maximum_memory_bytes, maximum_memory_bytes)	)
  resource.setrlimit(resource.RLIMIT_DATA						,	(maximum_memory_bytes, maximum_memory_bytes)	)
  if not platform.uname().system == "Darwin":
   resource.setrlimit(resource.RLIMIT_STACK						,	(maximum_memory_bytes, maximum_memory_bytes)	)

 faulthandler.disable()

 import builtins

 UpperCamelCase_     :					Any   	=None
 UpperCamelCase_     :					Any   	=None

 import os

 UpperCamelCase_     :					Union[str, Any]   	='1'

 UpperCamelCase_     :					Union[str, Any]   	=None
 UpperCamelCase_     :					Optional[Any]   	=None
 UpperCamelCase_     :					Optional[Any]   	=None
 UpperCamelCase_     :					Optional[Any]   	=None
 UpperCamelCase_     :					Dict   	=None
 UpperCamelCase_     :					List[Any]   	=None
 UpperCamelCase_     :					Any   	=None
 UpperCamelCase_     :					str   	=None
 UpperCamelCase_     :					Optional[int]   	=None
 UpperCamelCase_     :					Union[str, Any]   	=None
 UpperCamelCase_     :					str   	=None
 UpperCamelCase_     :					str   	=None
 UpperCamelCase_     :					List[Any]   	=None
 UpperCamelCase_     :					Optional[Any]   	=None
 UpperCamelCase_     :					str   	=None
 UpperCamelCase_     :					List[str]   	=None
 UpperCamelCase_     :					Optional[Any]   	=None
 UpperCamelCase_     :					int   	=None
 UpperCamelCase_     :					Optional[Any]   	=None
 UpperCamelCase_     :					Tuple   	=None
 UpperCamelCase_     :					List[Any]   	=None
 UpperCamelCase_     :					Tuple   	=None
 UpperCamelCase_     :					Dict   	=None
 UpperCamelCase_     :					Optional[int]   	=None
 UpperCamelCase_     :					str   	=None
 UpperCamelCase_     :					Union[str, Any]   	=None
 UpperCamelCase_     :					int   	=None

 import shutil

 UpperCamelCase_     :					List[str]   	=None
 UpperCamelCase_     :					Optional[Any]   	=None
 UpperCamelCase_     :					Dict   	=None

 import subprocess

 UpperCamelCase_     :					Optional[int]   	=None  # type: ignore

 UpperCamelCase_     :					Union[str, Any]   	=None

 import sys

 UpperCamelCase_     :					Union[str, Any]   	=None
 UpperCamelCase_     :					Tuple   	=None
 UpperCamelCase_     :					Union[str, Any]   	=None
 UpperCamelCase_     :					Dict   	=None
 UpperCamelCase_     :					Any   	=None
