


'''simple docstring'''






import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import (
		    MODEL_FOR_PRETRAINING_MAPPING,
		    MegatronBertForCausalLM,
		    MegatronBertForMaskedLM,
		    MegatronBertForMultipleChoice,
		    MegatronBertForNextSentencePrediction,
		    MegatronBertForPreTraining,
		    MegatronBertForQuestionAnswering,
		    MegatronBertForSequenceClassification,
		    MegatronBertForTokenClassification,
		    MegatronBertModel,
		)




class 						a__							:





							'''simple docstring'''






							def __init__(						self   ,       lowerCamelCase_   ,       lowerCamelCase_=13   ,       lowerCamelCase_=7   ,       lowerCamelCase_=True   ,       lowerCamelCase_=True   ,       lowerCamelCase_=True   ,       lowerCamelCase_=True   ,       lowerCamelCase_=99   ,       lowerCamelCase_=64   ,       lowerCamelCase_=32   ,       lowerCamelCase_=5   ,       lowerCamelCase_=4   ,       lowerCamelCase_=37   ,       lowerCamelCase_="gelu"   ,       lowerCamelCase_=0.1   ,       lowerCamelCase_=0.1   ,       lowerCamelCase_=5_12   ,       lowerCamelCase_=16   ,       lowerCamelCase_=2   ,       lowerCamelCase_=0.02   ,       lowerCamelCase_=3   ,       lowerCamelCase_=4   ,       lowerCamelCase_=None   ,       )			->				List[Any]:
														lowerCAmelCase__			  =      parent
														lowerCAmelCase__			  =      batch_size
														lowerCAmelCase__			  =      seq_length
														lowerCAmelCase__			  =      is_training
														lowerCAmelCase__			  =      use_input_mask
														lowerCAmelCase__			  =      use_token_type_ids
														lowerCAmelCase__			  =      use_labels
														lowerCAmelCase__			  =      vocab_size
														lowerCAmelCase__			  =      hidden_size
														lowerCAmelCase__			  =      embedding_size
														lowerCAmelCase__			  =      num_hidden_layers
														lowerCAmelCase__			  =      num_attention_heads
														lowerCAmelCase__			  =      intermediate_size
														lowerCAmelCase__			  =      hidden_act
														lowerCAmelCase__			  =      hidden_dropout_prob
														lowerCAmelCase__			  =      attention_probs_dropout_prob
														lowerCAmelCase__			  =      max_position_embeddings
														lowerCAmelCase__			  =      type_vocab_size
														lowerCAmelCase__			  =      type_sequence_label_size
														lowerCAmelCase__			  =      initializer_range
														lowerCAmelCase__			  =      num_labels
														lowerCAmelCase__			  =      num_choices
														lowerCAmelCase__			  =      scope







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														lowerCAmelCase__			  =      ids_tensor([self.batch_size, self.seq_length]   ,       self.vocab_size  )

														lowerCAmelCase__			  =      None
														if self.use_input_mask:
																					lowerCAmelCase__			  =      random_attention_mask([self.batch_size, self.seq_length]  )

														lowerCAmelCase__			  =      None
														if self.use_token_type_ids:
																					lowerCAmelCase__			  =      ids_tensor([self.batch_size, self.seq_length]   ,       self.type_vocab_size  )

														lowerCAmelCase__			  =      None
														lowerCAmelCase__			  =      None
														lowerCAmelCase__			  =      None
														if self.use_labels:
																					lowerCAmelCase__			  =      ids_tensor([self.batch_size]   ,       self.type_sequence_label_size  )
																					lowerCAmelCase__			  =      ids_tensor([self.batch_size, self.seq_length]   ,       self.num_labels  )
																					lowerCAmelCase__			  =      ids_tensor([self.batch_size]   ,       self.num_choices  )

														lowerCAmelCase__			  =      self.get_config()

														return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														return MegatronBertConfig(
														    vocab_size=self.vocab_size   ,       hidden_size=self.hidden_size   ,       num_hidden_layers=self.num_hidden_layers   ,       num_attention_heads=self.num_attention_heads   ,       intermediate_size=self.intermediate_size   ,       embedding_size=self.embedding_size   ,       hidden_act=self.hidden_act   ,       hidden_dropout_prob=self.hidden_dropout_prob   ,       attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,       max_position_embeddings=self.max_position_embeddings   ,       type_vocab_size=self.type_vocab_size   ,       is_decoder=lowerCamelCase_   ,       initializer_range=self.initializer_range   ,       )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				Any:
														lowerCAmelCase__			  =      MegatronBertModel(config=lowerCamelCase_  )
														model.to(lowerCamelCase_  )
														model.eval()
														lowerCAmelCase__			  =      model(lowerCamelCase_   ,       attention_mask=lowerCamelCase_   ,       token_type_ids=lowerCamelCase_  )
														lowerCAmelCase__			  =      model(lowerCamelCase_   ,       token_type_ids=lowerCamelCase_  )
														lowerCAmelCase__			  =      model(lowerCamelCase_  )

														self.parent.assertEqual(result.last_hidden_state.shape   ,       (self.batch_size, self.seq_length, self.hidden_size)  )
														self.parent.assertEqual(result.pooler_output.shape   ,       (self.batch_size, self.hidden_size)  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				Optional[Any]:
														lowerCAmelCase__			  =      MegatronBertForMaskedLM(config=lowerCamelCase_  )
														model.to(lowerCamelCase_  )
														model.eval()
														lowerCAmelCase__			  =      model(lowerCamelCase_   ,       attention_mask=lowerCamelCase_   ,       token_type_ids=lowerCamelCase_   ,       labels=lowerCamelCase_  )
														self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.seq_length, self.vocab_size)  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				Any:
														lowerCAmelCase__			  =      MegatronBertForCausalLM(config=lowerCamelCase_  )
														model.to(lowerCamelCase_  )
														model.eval()
														lowerCAmelCase__			  =      model(lowerCamelCase_   ,       attention_mask=lowerCamelCase_   ,       token_type_ids=lowerCamelCase_   ,       labels=lowerCamelCase_  )
														self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.seq_length, self.vocab_size)  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				int:
														lowerCAmelCase__			  =      MegatronBertForNextSentencePrediction(config=lowerCamelCase_  )
														model.to(lowerCamelCase_  )
														model.eval()
														lowerCAmelCase__			  =      model(
														    lowerCamelCase_   ,       attention_mask=lowerCamelCase_   ,       token_type_ids=lowerCamelCase_   ,       labels=lowerCamelCase_   ,       )
														self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, 2)  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				int:
														lowerCAmelCase__			  =      MegatronBertForPreTraining(config=lowerCamelCase_  )
														model.to(lowerCamelCase_  )
														model.eval()
														lowerCAmelCase__			  =      model(
														    lowerCamelCase_   ,       attention_mask=lowerCamelCase_   ,       token_type_ids=lowerCamelCase_   ,       labels=lowerCamelCase_   ,       next_sentence_label=lowerCamelCase_   ,       )
														self.parent.assertEqual(result.prediction_logits.shape   ,       (self.batch_size, self.seq_length, self.vocab_size)  )
														self.parent.assertEqual(result.seq_relationship_logits.shape   ,       (self.batch_size, 2)  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				Dict:
														lowerCAmelCase__			  =      MegatronBertForQuestionAnswering(config=lowerCamelCase_  )
														model.to(lowerCamelCase_  )
														model.eval()
														lowerCAmelCase__			  =      model(
														    lowerCamelCase_   ,       attention_mask=lowerCamelCase_   ,       token_type_ids=lowerCamelCase_   ,       start_positions=lowerCamelCase_   ,       end_positions=lowerCamelCase_   ,       )
														self.parent.assertEqual(result.start_logits.shape   ,       (self.batch_size, self.seq_length)  )
														self.parent.assertEqual(result.end_logits.shape   ,       (self.batch_size, self.seq_length)  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				List[str]:
														lowerCAmelCase__			  =      self.num_labels
														lowerCAmelCase__			  =      MegatronBertForSequenceClassification(lowerCamelCase_  )
														model.to(lowerCamelCase_  )
														model.eval()
														lowerCAmelCase__			  =      model(lowerCamelCase_   ,       attention_mask=lowerCamelCase_   ,       token_type_ids=lowerCamelCase_   ,       labels=lowerCamelCase_  )
														self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.num_labels)  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				Optional[int]:
														lowerCAmelCase__			  =      self.num_labels
														lowerCAmelCase__			  =      MegatronBertForTokenClassification(config=lowerCamelCase_  )
														model.to(lowerCamelCase_  )
														model.eval()
														lowerCAmelCase__			  =      model(lowerCamelCase_   ,       attention_mask=lowerCamelCase_   ,       token_type_ids=lowerCamelCase_   ,       labels=lowerCamelCase_  )
														self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.seq_length, self.num_labels)  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				Any:
														lowerCAmelCase__			  =      self.num_choices
														lowerCAmelCase__			  =      MegatronBertForMultipleChoice(config=lowerCamelCase_  )
														model.to(lowerCamelCase_  )
														model.eval()
														lowerCAmelCase__			  =      input_ids.unsqueeze(1  ).expand(-1   ,       self.num_choices   ,       -1  ).contiguous()
														lowerCAmelCase__			  =      token_type_ids.unsqueeze(1  ).expand(-1   ,       self.num_choices   ,       -1  ).contiguous()
														lowerCAmelCase__			  =      input_mask.unsqueeze(1  ).expand(-1   ,       self.num_choices   ,       -1  ).contiguous()
														lowerCAmelCase__			  =      model(
														    lowerCamelCase_   ,       attention_mask=lowerCamelCase_   ,       token_type_ids=lowerCamelCase_   ,       labels=lowerCamelCase_   ,       )
														self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.num_choices)  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														lowerCAmelCase__			  =      self.prepare_config_and_inputs()
														(
														    (
														    lowerCAmelCase__
														)  ,   (
														    lowerCAmelCase__
														)  ,   (
														    lowerCAmelCase__
														)  ,   (
														    lowerCAmelCase__
														)  ,   (
														    lowerCAmelCase__
														)  ,   (
														    lowerCAmelCase__
														)  ,   (
														    lowerCAmelCase__
														)  ,   
														)			  =      config_and_inputs
														lowerCAmelCase__			  =      {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
														return config, inputs_dict




@require_torch
class 						a__							(				a__							,							a__							,							unittest.TestCase    ):





							'''simple docstring'''






							lowercase__     :  Any      		= (
							    (
							        MegatronBertModel,
							        MegatronBertForMaskedLM,
							        MegatronBertForCausalLM,
							        MegatronBertForMultipleChoice,
							        MegatronBertForNextSentencePrediction,
							        MegatronBertForPreTraining,
							        MegatronBertForQuestionAnswering,
							        MegatronBertForSequenceClassification,
							        MegatronBertForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							lowercase__     :  Union[str, Any]      		= (
							    {
							        "feature-extraction": MegatronBertModel,
							        "fill-mask": MegatronBertForMaskedLM,
							        "question-answering": MegatronBertForQuestionAnswering,
							        "text-classification": MegatronBertForSequenceClassification,
							        "text-generation": MegatronBertForCausalLM,
							        "token-classification": MegatronBertForTokenClassification,
							        "zero-shot": MegatronBertForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							lowercase__     :  Dict      		= True
							# test_resize_embeddings = False
							lowercase__     :  Optional[Any]      		= False







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_=False  )			->				Optional[int]:
														lowerCAmelCase__			  =      super()._prepare_for_class(lowerCamelCase_   ,       lowerCamelCase_   ,       return_labels=lowerCamelCase_  )

														if return_labels:
																					if model_class in get_values(lowerCamelCase_  ):
																												lowerCAmelCase__			  =      torch.zeros(
																												    (self.model_tester.batch_size, self.model_tester.seq_length)   ,       dtype=torch.long   ,       device=lowerCamelCase_  )
																												lowerCAmelCase__			  =      torch.zeros(
																												    self.model_tester.batch_size   ,       dtype=torch.long   ,       device=lowerCamelCase_  )
														return inputs_dict







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[str]:
														lowerCAmelCase__			  =      MegatronBertModelTester(self  )
														lowerCAmelCase__			  =      ConfigTester(self   ,       config_class=lowerCamelCase_   ,       hidden_size=37  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[str]:
														self.config_tester.run_common_tests()







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														lowerCAmelCase__			  =      self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_model(*lowerCamelCase_  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[str]:
														lowerCAmelCase__			  =      self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowerCamelCase_  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[str]:
														lowerCAmelCase__			  =      self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCamelCase_  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														lowerCAmelCase__			  =      self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowerCamelCase_  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[Any]:
														lowerCAmelCase__			  =      self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCamelCase_  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														lowerCAmelCase__			  =      self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCamelCase_  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														lowerCAmelCase__			  =      self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCamelCase_  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[Any]:
														lowerCAmelCase__			  =      self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowerCamelCase_  )


def    _snake_case  (	A			)		->       Any:
							return torch.tensor(
							    A					,       dtype=torch.long					,       device=A					,       )


__UpperCAmelCase       							=				1e-4






@require_torch
@require_sentencepiece
@require_tokenizers
class 						a__							(				unittest.TestCase    ):





							'''simple docstring'''






							@slow
							@unittest.skip('''Model is not available.'''  )
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[Any]:
														lowerCAmelCase__			  =      '''nvidia/megatron-bert-uncased-345m'''
														if "MYDIR" in os.environ:
																					lowerCAmelCase__			  =      os.path.join(os.environ['''MYDIR''']   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      MegatronBertModel.from_pretrained(lowerCamelCase_  )
														model.to(lowerCamelCase_  )
														model.half()
														lowerCAmelCase__			  =      _long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]  )
														with torch.no_grad():
																					lowerCAmelCase__			  =      model(lowerCamelCase_  )[0]
														lowerCAmelCase__			  =      torch.Size((1, 9, 10_24)  )
														self.assertEqual(output.shape   ,       lowerCamelCase_  )

														lowerCAmelCase__			  =      [-0.6_040, -0.2_517, -0.1_025, 0.3_420, -0.6_758, -0.0_017, -0.1_089, -0.1_990, 0.5_728]
														for ii in range(3  ):
																					for jj in range(3  ):
																												lowerCAmelCase__			  =      output[0, ii, jj]
																												lowerCAmelCase__			  =      expected[3 * ii + jj]
																												lowerCAmelCase__			  =      '''ii={} jj={} a={} b={}'''.format(lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )
																												self.assertTrue(math.isclose(lowerCamelCase_   ,       lowerCamelCase_   ,       rel_tol=lowerCamelCase_   ,       abs_tol=lowerCamelCase_  )   ,       msg=lowerCamelCase_  )






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


_lowerCAmelCase				:   Dict						       =						logging.getLogger(__name__)


@dataclass
class        __magic_name__	:





     """simple docstring"""



     __UpperCamelCase            =  field(
         metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models'''}				)
     __UpperCamelCase            =  field(
         default=lowerCamelCase__     ,					metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co'''}     ,					)
     __UpperCamelCase            =  field(
         default=lowerCamelCase__     ,					metadata={'''help''': '''Whether to freeze the feature extractor layers of the model.'''}				)
     __UpperCamelCase            =  field(
         default=lowerCamelCase__     ,					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.99_9995     ,					metadata={'''help''': '''Decay of gumbel temperature during training.'''}				)




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


@dataclass
class        __magic_name__	:





     """simple docstring"""



     __UpperCamelCase            =  field(
         default=lowerCamelCase__     ,					metadata={'''help''': '''The name of the dataset to use (via the datasets library).'''}				)
     __UpperCamelCase            =  field(
         default=lowerCamelCase__     ,					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=lowerCamelCase__     ,					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=lowerCamelCase__     ,					metadata={'''help''': '''The number of processes to use for the preprocessing.'''}     ,					)
     __UpperCamelCase            =  field(
         default=20.0     ,					metadata={'''help''': '''Filter audio files that are longer than `max_duration_in_seconds` seconds'''}				)




@dataclass
class        __magic_name__	:





     """simple docstring"""



     __UpperCamelCase            =  42
     __UpperCamelCase            =  42
     __UpperCamelCase            =  "longest"
     __UpperCamelCase            =  None
     __UpperCamelCase            =  None



     def __call__(				self							:Optional[Any]      ,   snake_case							:List[Dict[str, Union[List[int], torch.Tensor]]]	):





        '''simple docstring'''




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

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

           A_       :					Union[str, 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
           A_       :					str        =      1
           A_       :					Union[str, Any]        =      attention_mask.flip([-1]	).cumsum(-1	).flip([-1]	).bool()

        # sample randomly masked indices
        A_       :					Optional[int]        =      _compute_mask_indices(
            (batch_size, mask_indices_seq_length)      ,   self.model.config.mask_time_prob      ,   self.model.config.mask_time_length      ,   attention_mask=snake_case      ,   min_masks=2      ,   )

        return batch




class        __magic_name__	(				lowerCamelCase__				):





     """simple docstring"""



     def __init__(				self							:int      ,   *snake_case							:Any      ,   snake_case							:Any=1      ,   snake_case							:Dict=0      ,   snake_case							:Dict=1.0      ,   **snake_case							:Any	):





        '''simple docstring'''




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





     def     SCREAMING_SNAKE_CASE		(				self							:Optional[int]      ,   snake_case							:nn.Module      ,   snake_case							:Dict[str, Union[torch.Tensor, Any]]	):





        '''simple docstring'''




        model.train()
        A_       :					List[str]        =      self._prepare_inputs(snake_case	)

        if self.use_amp:
           with autocast():
              A_       :					List[str]        =      self.compute_loss(snake_case      ,   snake_case	)
        else:
           A_       :					str        =      self.compute_loss(snake_case      ,   snake_case	)

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

        if self.use_amp:
           self.scaler.scale(snake_case	).backward()
        elif self.use_apex:
           with amp.scale_loss(snake_case      ,   self.optimizer	) as scaled_loss:
              scaled_loss.backward()
        elif self.deepspeed:
           self.deepspeed.backward(snake_case	)
        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   __snake_case    (       )			->    Tuple:
   # 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.

   A_       :					List[str]        =      HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)						)

   A_   ,    A_   ,    A_       :					Dict        =      parser.parse_args_into_dataclasses()
   configure_logger(_lowerCAmelCase					,					_lowerCAmelCase						)

   # Downloading and loading a dataset from the hub.
   A_       :					Optional[Any]        =      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_       :					Tuple        =      DatasetDict()
      A_       :					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					,					)
      A_       :					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					,					)
   else:
      # make sure only "validation" and "train" keys remain"
      A_       :					Any        =      DatasetDict()
      A_       :					Optional[Any]        =      load_dataset(
          data_args.dataset_name					,					data_args.dataset_config_name					,					split="validation"					,					cache_dir=model_args.cache_dir					,					)
      A_       :					List[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
   A_       :					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			:			str						):
      # check that all files have the correct sampling rate
      A_   ,    A_       :					Tuple        =      librosa.load(batch[data_args.speech_file_column]					,					sr=feature_extractor.sampling_rate						)
      return batch

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

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

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

   # normalize and transform to `BatchFeatures`
   A_       :					Any        =      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
   A_       :					Any        =      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_       :					int        =      WavaVecaForPreTraining(_lowerCAmelCase						)

   A_       :					str        =      DataCollatorForWavaVecaPretraining(model=_lowerCAmelCase					,					feature_extractor=_lowerCAmelCase						)

   A_       :					List[Any]        =      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()
