

"""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 lowerCamelCase__   :



    '''simple docstring'''




    def __init__( self ,lowerCamelCase_ ,lowerCamelCase_=1_3 ,lowerCamelCase_=7 ,lowerCamelCase_=True ,lowerCamelCase_=True ,lowerCamelCase_=True ,lowerCamelCase_=True ,lowerCamelCase_=9_9 ,lowerCamelCase_=6_4 ,lowerCamelCase_=3_2 ,lowerCamelCase_=5 ,lowerCamelCase_=4 ,lowerCamelCase_=3_7 ,lowerCamelCase_="gelu" ,lowerCamelCase_=0.1 ,lowerCamelCase_=0.1 ,lowerCamelCase_=5_1_2 ,lowerCamelCase_=1_6 ,lowerCamelCase_=2 ,lowerCamelCase_=0.02 ,lowerCamelCase_=3 ,lowerCamelCase_=4 ,lowerCamelCase_=None ,)    ->      List[Any]:
          A				  =		parent
          A				  =		batch_size
          A				  =		seq_length
          A				  =		is_training
          A				  =		use_input_mask
          A				  =		use_token_type_ids
          A				  =		use_labels
          A				  =		vocab_size
          A				  =		hidden_size
          A				  =		embedding_size
          A				  =		num_hidden_layers
          A				  =		num_attention_heads
          A				  =		intermediate_size
          A				  =		hidden_act
          A				  =		hidden_dropout_prob
          A				  =		attention_probs_dropout_prob
          A				  =		max_position_embeddings
          A				  =		type_vocab_size
          A				  =		type_sequence_label_size
          A				  =		initializer_range
          A				  =		num_labels
          A				  =		num_choices
          A				  =		scope




    def UpperCamelCase__		( self    )    ->      Optional[Any]:
          A				  =		ids_tensor([self.batch_size, self.seq_length] ,self.vocab_size    )

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

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

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

          A				  =		self.get_config()

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




    def UpperCamelCase__		( self    )    ->      Dict:
          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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Tuple:
          A				  =		MegatronBertModel(config=lowerCamelCase_    )
          model.to(lowerCamelCase_    )
          model.eval()
          A				  =		model(lowerCamelCase_ ,attention_mask=lowerCamelCase_ ,token_type_ids=lowerCamelCase_    )
          A				  =		model(lowerCamelCase_ ,token_type_ids=lowerCamelCase_    )
          A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Optional[int]:
          A				  =		MegatronBertForMaskedLM(config=lowerCamelCase_    )
          model.to(lowerCamelCase_    )
          model.eval()
          A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Union[str, Any]:
          A				  =		MegatronBertForCausalLM(config=lowerCamelCase_    )
          model.to(lowerCamelCase_    )
          model.eval()
          A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Dict:
          A				  =		MegatronBertForNextSentencePrediction(config=lowerCamelCase_    )
          model.to(lowerCamelCase_    )
          model.eval()
          A				  =		model(
              lowerCamelCase_ ,attention_mask=lowerCamelCase_ ,token_type_ids=lowerCamelCase_ ,labels=lowerCamelCase_ ,)
          self.parent.assertEqual(result.logits.shape ,(self.batch_size, 2)    )




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Optional[Any]:
          A				  =		MegatronBertForPreTraining(config=lowerCamelCase_    )
          model.to(lowerCamelCase_    )
          model.eval()
          A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Dict:
          A				  =		MegatronBertForQuestionAnswering(config=lowerCamelCase_    )
          model.to(lowerCamelCase_    )
          model.eval()
          A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      str:
          A				  =		self.num_labels
          A				  =		MegatronBertForSequenceClassification(lowerCamelCase_    )
          model.to(lowerCamelCase_    )
          model.eval()
          A				  =		model(lowerCamelCase_ ,attention_mask=lowerCamelCase_ ,token_type_ids=lowerCamelCase_ ,labels=lowerCamelCase_    )
          self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.num_labels)    )




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Optional[Any]:
          A				  =		self.num_labels
          A				  =		MegatronBertForTokenClassification(config=lowerCamelCase_    )
          model.to(lowerCamelCase_    )
          model.eval()
          A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Dict:
          A				  =		self.num_choices
          A				  =		MegatronBertForMultipleChoice(config=lowerCamelCase_    )
          model.to(lowerCamelCase_    )
          model.eval()
          A				  =		input_ids.unsqueeze(1    ).expand(-1 ,self.num_choices ,-1    ).contiguous()
          A				  =		token_type_ids.unsqueeze(1    ).expand(-1 ,self.num_choices ,-1    ).contiguous()
          A				  =		input_mask.unsqueeze(1    ).expand(-1 ,self.num_choices ,-1    ).contiguous()
          A				  =		model(
              lowerCamelCase_ ,attention_mask=lowerCamelCase_ ,token_type_ids=lowerCamelCase_ ,labels=lowerCamelCase_ ,)
          self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.num_choices)    )




    def UpperCamelCase__		( self    )    ->      Union[str, Any]:
          A				  =		self.prepare_config_and_inputs()
          (
              (
              A
          )  ,			(
              A
          )  ,			(
              A
          )  ,			(
              A
          )  ,			(
              A
          )  ,			(
              A
          )  ,			(
              A
          )  ,			
          )				  =		config_and_inputs
          A				  =		{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
          return config, inputs_dict





@require_torch
class lowerCamelCase__   (				SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE     ,   unittest.TestCase    ):



    '''simple docstring'''




    _lowerCamelCase					=   (
        (
            MegatronBertModel,
            MegatronBertForMaskedLM,
            MegatronBertForCausalLM,
            MegatronBertForMultipleChoice,
            MegatronBertForNextSentencePrediction,
            MegatronBertForPreTraining,
            MegatronBertForQuestionAnswering,
            MegatronBertForSequenceClassification,
            MegatronBertForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    _lowerCamelCase					=   (
        {
            '''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 {}
    )
    _lowerCamelCase					=   True
    # test_resize_embeddings = False
    _lowerCamelCase					=   False




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_=False    )    ->      str:
          A				  =		super()._prepare_for_class(lowerCamelCase_ ,lowerCamelCase_ ,return_labels=lowerCamelCase_    )

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




    def UpperCamelCase__		( self    )    ->      str:
          A				  =		MegatronBertModelTester(self    )
          A				  =		ConfigTester(self ,config_class=lowerCamelCase_ ,hidden_size=3_7    )




    def UpperCamelCase__		( self    )    ->      Optional[Any]:
          self.config_tester.run_common_tests()




    def UpperCamelCase__		( self    )    ->      str:
          A				  =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_model(*lowerCamelCase_    )




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




    def UpperCamelCase__		( self    )    ->      Tuple:
          A				  =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCamelCase_    )




    def UpperCamelCase__		( self    )    ->      List[Any]:
          A				  =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowerCamelCase_    )




    def UpperCamelCase__		( self    )    ->      str:
          A				  =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCamelCase_    )




    def UpperCamelCase__		( self    )    ->      str:
          A				  =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCamelCase_    )




    def UpperCamelCase__		( self    )    ->      str:
          A				  =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCamelCase_    )




    def UpperCamelCase__		( self    )    ->      Any:
          A				  =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowerCamelCase_    )




def   _A					(	_a	:						List[str]				):



      """simple docstring"""






      return torch.tensor(
          _a     ,   dtype=torch.long     ,   device=_a     ,   )


UpperCAmelCase													=1E-4





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



    '''simple docstring'''




    @slow
    @unittest.skip("""Model is not available."""    )
    def UpperCamelCase__		( self    )    ->      List[str]:
          A				  =		"""nvidia/megatron-bert-uncased-345m"""
          if "MYDIR" in os.environ:
                A				  =		os.path.join(os.environ["""MYDIR"""] ,lowerCamelCase_    )
          A				  =		MegatronBertModel.from_pretrained(lowerCamelCase_    )
          model.to(lowerCamelCase_    )
          model.half()
          A				  =		_long_tensor([[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]]    )
          with torch.no_grad():
                A				  =		model(lowerCamelCase_    )[0]
          A				  =		torch.Size((1, 9, 1_0_2_4)    )
          self.assertEqual(output.shape ,lowerCamelCase_    )

          A				  =		[-0.60_40, -0.25_17, -0.10_25, 0.34_20, -0.67_58, -0.00_17, -0.10_89, -0.19_90, 0.57_28]
          for ii in range(3    ):
                for jj in range(3    ):
                      A				  =		output[0, ii, jj]
                      A				  =		expected[3 * ii + jj]
                      A				  =		"""ii={} jj={} a={} b={}""".format(lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )
                      self.assertTrue(math.isclose(lowerCamelCase_ ,lowerCamelCase_ ,rel_tol=lowerCamelCase_ ,abs_tol=lowerCamelCase_    ) ,msg=lowerCamelCase_    )







"""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 lowerCamelCase__   :



    '''simple docstring'''




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




def   _A					(	_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				)]     ,   )
      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(_a				)





@dataclass
class lowerCamelCase__   :



    '''simple docstring'''




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





@dataclass
class lowerCamelCase__   :



    '''simple docstring'''




    _lowerCamelCase					=   42
    _lowerCamelCase					=   42
    _lowerCamelCase					=   "longest"
    _lowerCamelCase					=   None
    _lowerCamelCase					=   None




    def __call__( self ,lowerCamelCase_    )    ->      Dict[str, torch.Tensor]:
          # reformat list to dict and set to pytorch format
          A				  =		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""" ,)
          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=lowerCamelCase_ ,min_masks=2 ,)

          return batch





class lowerCamelCase__   (				SCREAMING_SNAKE_CASE    ):



    '''simple docstring'''




    def __init__( self ,*lowerCamelCase_ ,lowerCamelCase_=1 ,lowerCamelCase_=0 ,lowerCamelCase_=1.0 ,**lowerCamelCase_    )    ->      Union[str, Any]:
          super().__init__(*lowerCamelCase_ ,**lowerCamelCase_    )
          A				  =		0
          A				  =		max_gumbel_temp
          A				  =		min_gumbel_temp
          A				  =		gumbel_temp_decay




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_    )    ->      torch.Tensor:
          model.train()
          A				  =		self._prepare_inputs(lowerCamelCase_    )

          if self.use_amp:
                with autocast():
                      A				  =		self.compute_loss(lowerCamelCase_ ,lowerCamelCase_    )
          else:
                A				  =		self.compute_loss(lowerCamelCase_ ,lowerCamelCase_    )

          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(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   _A					(	):



      """simple docstring"""






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

      A  ,			A  ,			A				  =		parser.parse_args_into_dataclasses()
      configure_logger(_a     ,   _a				)

      # 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=_a				)

      def prepare_dataset(_a	:						Dict				):
            # 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(
          _a     ,   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 _a				: len(data["""speech"""]				) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate				)				)

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

      # normalize and transform to `BatchFeatures`
      A				  =		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
      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(_a				)

      A				  =		DataCollatorForWavaVecaPretraining(model=_a     ,   feature_extractor=_a				)

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




