
"""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 				snake_case    :
   a_ :						str			 =			field(
       metadata={"""help""": """Path to pretrained model or model identifier from huggingface.co/models"""}						)
   a_ :						Optional[str]			 =			field(
       default=__A				,    metadata={"""help""": """Where do you want to store the pretrained models downloaded from huggingface.co"""}				,    )
   a_ :						Optional[bool]			 =			field(
       default=__A				,    metadata={"""help""": """Whether to freeze the feature extractor layers of the model."""}						)
   a_ :						Optional[bool]			 =			field(
       default=__A				,    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.999995				,    metadata={"""help""": """Decay of gumbel temperature during training."""}						)






def 					UpperCamelCase		( UpperCAmelCase      ,					UpperCAmelCase			)		->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			)]      ,					)
   a_							=			logging.WARNING
   if model_args.verbose_logging:
      a_							=			logging.DEBUG
   elif trainer_utils.is_main_process(training_args.local_rank			):
      a_							=			logging.INFO
   logger.setLevel(SCREAMING_SNAKE_CASE__			)





@dataclass
class 				snake_case    :
   a_ :						str			 =			field(
       default=__A				,    metadata={"""help""": """The name of the dataset to use (via the datasets library)."""}						)
   a_ :						Optional[str]			 =			field(
       default=__A				,    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=__A				,    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=__A				,    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 				snake_case    :
   a_ :						WavaVecaForPreTraining
   a_ :						WavaVecaFeatureExtractor
   a_ :						Union[bool, str]			 =			"longest"
   a_ :						Optional[int]			 =			None
   a_ :						Optional[int]			 =			None







   def __call__(      self    ,				__UpperCAmelCase)					->Dict[str, torch.Tensor]:
      # reformat list to dict and set to pytorch format
      a_							=			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"    ,				)
      a_							=			self.model._get_feat_extract_output_lengths(batch["input_values"].shape[-1])

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

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

      # sample randomly masked indices
      a_							=			_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 				snake_case    (			__A						):







   def __init__(      self    ,				*__UpperCAmelCase    ,				__UpperCAmelCase=1    ,				__UpperCAmelCase=0    ,				__UpperCAmelCase=1.0    ,				**__UpperCAmelCase)					->str:
      super().__init__(*_UpperCamelCase    ,				**_UpperCamelCase)
      a_							=			0
      a_							=			max_gumbel_temp
      a_							=			min_gumbel_temp
      a_							=			gumbel_temp_decay







   def 					UpperCAmelCase__     (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase)					->torch.Tensor:
      model.train()
      a_							=			self._prepare_inputs(_UpperCamelCase)

      if self.use_amp:
         with autocast():
            a_							=			self.compute_loss(_UpperCamelCase    ,				_UpperCamelCase)
      else:
         a_							=			self.compute_loss(_UpperCamelCase    ,				_UpperCamelCase)

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

   """simple docstring"""



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

   a_					,       a_					,       a_							=			parser.parse_args_into_dataclasses()
   configure_logger(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__			)

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

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

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

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

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

   # normalize and transform to `BatchFeatures`
   a_							=			vectorized_datasets.map(
       SCREAMING_SNAKE_CASE__      ,					batched=SCREAMING_SNAKE_CASE__      ,					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
   a_							=			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\'"			)

   a_							=			WavaVecaForPreTraining(SCREAMING_SNAKE_CASE__			)

   a_							=			DataCollatorForWavaVecaPretraining(model=SCREAMING_SNAKE_CASE__      ,					feature_extractor=SCREAMING_SNAKE_CASE__			)

   a_							=			WavaVecaPreTrainer(
       model=SCREAMING_SNAKE_CASE__      ,					data_collator=SCREAMING_SNAKE_CASE__      ,					args=SCREAMING_SNAKE_CASE__      ,					train_dataset=vectorized_datasets["train"]      ,					eval_dataset=vectorized_datasets["validation"]      ,					tokenizer=SCREAMING_SNAKE_CASE__      ,					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()


import unittest

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate


def   __SCREAMING_SNAKE_CASE					(SCREAMING_SNAKE_CASE__       ):
						return x + 2





class 	snake_case_   (		unittest.TestCase					):







			'''simple docstring'''

			def 				snake_case__(     self :				Optional[Any]							)      ->int:
									snake_case_   =	'''x = 3'''
									snake_case_   =	{}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									assert result == 3
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3}							)

									snake_case_   =	'''x = y'''
									snake_case_   =	{'''y''': 5}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									# evaluate returns the value of the last assignment.
									assert result == 5
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 5, '''y''': 5}							)







			def 				snake_case__(     self :				Dict							)      ->Optional[int]:
									snake_case_   =	'''y = add_two(x)'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {'''add_two''': add_two}							,   state=_UpperCamelCase							)
									assert result == 5
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''y''': 5}							)

									# Won't work without the tool
									with CaptureStdout() as out:
															snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									assert result is None
									assert "tried to execute add_two" in out.out







			def 				snake_case__(     self :				Union[str, Any]							)      ->Dict:
									snake_case_   =	'''x = 3'''
									snake_case_   =	{}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									assert result == 3
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3}							)







			def 				snake_case__(     self :				Optional[int]							)      ->Optional[int]:
									snake_case_   =	'''test_dict = {\'x\': x, \'y\': add_two(x)}'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {'''add_two''': add_two}							,   state=_UpperCamelCase							)
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''y''': 5}							)
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''test_dict''': {'''x''': 3, '''y''': 5}}							)







			def 				snake_case__(     self :				Dict							)      ->str:
									snake_case_   =	'''x = 3\ny = 5'''
									snake_case_   =	{}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									# evaluate returns the value of the last assignment.
									assert result == 5
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''y''': 5}							)







			def 				snake_case__(     self :				str							)      ->Tuple:
									snake_case_   =	'''text = f\'This is x: {x}.\''''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									# evaluate returns the value of the last assignment.
									assert result == "This is x: 3."
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''text''': '''This is x: 3.'''}							)







			def 				snake_case__(     self :				Optional[Any]							)      ->List[str]:
									snake_case_   =	'''if x <= 3:\n    y = 2\nelse:\n    y = 5'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									# evaluate returns the value of the last assignment.
									assert result == 2
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''y''': 2}							)

									snake_case_   =	{'''x''': 8}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									# evaluate returns the value of the last assignment.
									assert result == 5
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 8, '''y''': 5}							)







			def 				snake_case__(     self :				str							)      ->str:
									snake_case_   =	'''test_list = [x, add_two(x)]'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {'''add_two''': add_two}							,   state=_UpperCamelCase							)
									self.assertListEqual(_UpperCamelCase							,   [3, 5]							)
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''test_list''': [3, 5]}							)







			def 				snake_case__(     self :				Any							)      ->List[Any]:
									snake_case_   =	'''y = x'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									assert result == 3
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''y''': 3}							)







			def 				snake_case__(     self :				Optional[int]							)      ->Dict:
									snake_case_   =	'''test_list = [x, add_two(x)]\ntest_list[1]'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {'''add_two''': add_two}							,   state=_UpperCamelCase							)
									assert result == 5
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''test_list''': [3, 5]}							)

									snake_case_   =	'''test_dict = {\'x\': x, \'y\': add_two(x)}\ntest_dict[\'y\']'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {'''add_two''': add_two}							,   state=_UpperCamelCase							)
									assert result == 5
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''test_dict''': {'''x''': 3, '''y''': 5}}							)







			def 				snake_case__(     self :				Optional[Any]							)      ->int:
									snake_case_   =	'''x = 0\nfor i in range(3):\n    x = i'''
									snake_case_   =	{}
									snake_case_   =	evaluate(_UpperCamelCase							,   {'''range''': range}							,   state=_UpperCamelCase							)
									assert result == 2
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 2, '''i''': 2}							)