






'''simple docstring'''





import collections
import inspect
import unittest

from transformers import FocalNetConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import (
     FocalNetBackbone,
     FocalNetForImageClassification,
     FocalNetForMaskedImageModeling,
     FocalNetModel,
 )
 from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor
class 				lowercase       :


    '''simple docstring'''

    def __init__(					self     :     int				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     Optional[int]=13				,     __lowerCamelCase     :     Tuple=32				,     __lowerCamelCase     :     str=2				,     __lowerCamelCase     :     Tuple=3				,     __lowerCamelCase     :     str=16				,     __lowerCamelCase     :     Union[str, Any]=[32, 64, 128]				,     __lowerCamelCase     :     Dict=[1, 2, 1]				,     __lowerCamelCase     :     List[Any]=[2, 2, 4]				,     __lowerCamelCase     :     Dict=2				,     __lowerCamelCase     :     List[Any]=2.0				,     __lowerCamelCase     :     Dict=True				,     __lowerCamelCase     :     Any=0.0				,     __lowerCamelCase     :     int=0.0				,     __lowerCamelCase     :     Optional[Any]=0.1				,     __lowerCamelCase     :     List[Any]="gelu"				,     __lowerCamelCase     :     str=False				,     __lowerCamelCase     :     Optional[int]=True				,     __lowerCamelCase     :     str=0.0_2				,     __lowerCamelCase     :     Tuple=1E-5				,     __lowerCamelCase     :     Tuple=True				,     __lowerCamelCase     :     Any=None				,     __lowerCamelCase     :     Optional[int]=True				,     __lowerCamelCase     :     Optional[Any]=10				,     __lowerCamelCase     :     Tuple=8				,     __lowerCamelCase     :     str=["stage1", "stage2"]				,     __lowerCamelCase     :     Optional[int]=[1, 2]				,     )							->							List[str]:

     '''simple docstring'''




     lowerCamelCase__  =     parent
     lowerCamelCase__  =     batch_size
     lowerCamelCase__  =     image_size
     lowerCamelCase__  =     patch_size
     lowerCamelCase__  =     num_channels
     lowerCamelCase__  =     embed_dim
     lowerCamelCase__  =     hidden_sizes
     lowerCamelCase__  =     depths
     lowerCamelCase__  =     num_heads
     lowerCamelCase__  =     window_size
     lowerCamelCase__  =     mlp_ratio
     lowerCamelCase__  =     qkv_bias
     lowerCamelCase__  =     hidden_dropout_prob
     lowerCamelCase__  =     attention_probs_dropout_prob
     lowerCamelCase__  =     drop_path_rate
     lowerCamelCase__  =     hidden_act
     lowerCamelCase__  =     use_absolute_embeddings
     lowerCamelCase__  =     patch_norm
     lowerCamelCase__  =     layer_norm_eps
     lowerCamelCase__  =     initializer_range
     lowerCamelCase__  =     is_training
     lowerCamelCase__  =     scope
     lowerCamelCase__  =     use_labels
     lowerCamelCase__  =     type_sequence_label_size
     lowerCamelCase__  =     encoder_stride
     lowerCamelCase__  =     out_features
     lowerCamelCase__  =     out_indices

    def    a__				(					self     :     Tuple     )							->							List[Any]:

     '''simple docstring'''




     lowerCamelCase__  =     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

     lowerCamelCase__  =     None
     if self.use_labels:
      lowerCamelCase__  =     ids_tensor([self.batch_size]				,     self.type_sequence_label_size     )

     lowerCamelCase__  =     self.get_config()

     return config, pixel_values, labels

    def    a__				(					self     :     Dict     )							->							Optional[Any]:

     '''simple docstring'''




     return FocalNetConfig(
         image_size=self.image_size				,     patch_size=self.patch_size				,     num_channels=self.num_channels				,     embed_dim=self.embed_dim				,     hidden_sizes=self.hidden_sizes				,     depths=self.depths				,     num_heads=self.num_heads				,     window_size=self.window_size				,     mlp_ratio=self.mlp_ratio				,     qkv_bias=self.qkv_bias				,     hidden_dropout_prob=self.hidden_dropout_prob				,     attention_probs_dropout_prob=self.attention_probs_dropout_prob				,     drop_path_rate=self.drop_path_rate				,     hidden_act=self.hidden_act				,     use_absolute_embeddings=self.use_absolute_embeddings				,     path_norm=self.patch_norm				,     layer_norm_eps=self.layer_norm_eps				,     initializer_range=self.initializer_range				,     encoder_stride=self.encoder_stride				,     out_features=self.out_features				,     out_indices=self.out_indices				,     )

    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     Any     )							->							Optional[int]:

     '''simple docstring'''




     lowerCamelCase__  =     FocalNetModel(config=__lowerCamelCase     )
     model.to(__lowerCamelCase     )
     model.eval()
     lowerCamelCase__  =     model(__lowerCamelCase     )

     lowerCamelCase__  =     ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths     ) - 1))
     lowerCamelCase__  =     int(config.embed_dim * 2 ** (len(config.depths     ) - 1)     )

     self.parent.assertEqual(result.last_hidden_state.shape				,     (self.batch_size, expected_seq_len, expected_dim)     )

    def    a__				(					self     :     Optional[Any]				,     __lowerCamelCase     :     List[str]				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     Optional[int]     )							->							str:

     '''simple docstring'''




     lowerCamelCase__  =     FocalNetBackbone(config=__lowerCamelCase     )
     model.to(__lowerCamelCase     )
     model.eval()
     lowerCamelCase__  =     model(__lowerCamelCase     )

     # verify feature maps
     self.parent.assertEqual(len(result.feature_maps     )				,     len(config.out_features     )     )
     self.parent.assertListEqual(list(result.feature_maps[0].shape     )				,     [self.batch_size, self.image_size, 8, 8]     )

     # verify channels
     self.parent.assertEqual(len(model.channels     )				,     len(config.out_features     )     )
     self.parent.assertListEqual(model.channels				,     config.hidden_sizes[:-1]     )

     # verify backbone works with out_features=None
     lowerCamelCase__  =     None
     lowerCamelCase__  =     FocalNetBackbone(config=__lowerCamelCase     )
     model.to(__lowerCamelCase     )
     model.eval()
     lowerCamelCase__  =     model(__lowerCamelCase     )

     # verify feature maps
     self.parent.assertEqual(len(result.feature_maps     )				,     1     )
     self.parent.assertListEqual(list(result.feature_maps[0].shape     )				,     [self.batch_size, self.image_size * 2, 4, 4]     )

     # verify channels
     self.parent.assertEqual(len(model.channels     )				,     1     )
     self.parent.assertListEqual(model.channels				,     [config.hidden_sizes[-1]]     )

    def    a__				(					self     :     Any				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     List[Any]     )							->							Union[str, Any]:

     '''simple docstring'''




     lowerCamelCase__  =     FocalNetForMaskedImageModeling(config=__lowerCamelCase     )
     model.to(__lowerCamelCase     )
     model.eval()
     lowerCamelCase__  =     model(__lowerCamelCase     )
     self.parent.assertEqual(
         result.reconstruction.shape				,     (self.batch_size, self.num_channels, self.image_size, self.image_size)     )

     # test greyscale images
     lowerCamelCase__  =     1
     lowerCamelCase__  =     FocalNetForMaskedImageModeling(__lowerCamelCase     )
     model.to(__lowerCamelCase     )
     model.eval()

     lowerCamelCase__  =     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
     lowerCamelCase__  =     model(__lowerCamelCase     )
     self.parent.assertEqual(result.reconstruction.shape				,     (self.batch_size, 1, self.image_size, self.image_size)     )

    def    a__				(					self     :     Union[str, Any]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     Dict     )							->							Union[str, Any]:

     '''simple docstring'''




     lowerCamelCase__  =     self.type_sequence_label_size
     lowerCamelCase__  =     FocalNetForImageClassification(__lowerCamelCase     )
     model.to(__lowerCamelCase     )
     model.eval()
     lowerCamelCase__  =     model(__lowerCamelCase				,     labels=__lowerCamelCase     )
     self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.type_sequence_label_size)     )

     # test greyscale images
     lowerCamelCase__  =     1
     lowerCamelCase__  =     FocalNetForImageClassification(__lowerCamelCase     )
     model.to(__lowerCamelCase     )
     model.eval()

     lowerCamelCase__  =     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
     lowerCamelCase__  =     model(__lowerCamelCase     )
     self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.type_sequence_label_size)     )




    def    a__				(					self     :     List[Any]     )							->							int:

     '''simple docstring'''




     lowerCamelCase__  =     self.prepare_config_and_inputs()

     lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  =     config_and_inputs
     lowerCamelCase__  =     {"pixel_values": pixel_values}
     return config, inputs_dict
@require_torch
class 				lowercase       (							_lowerCamelCase ,       _lowerCamelCase ,       unittest.TestCase      ):


    '''simple docstring'''

    lowerCAmelCase__      				=							(
        (
            FocalNetModel,
            FocalNetForImageClassification,
            FocalNetForMaskedImageModeling,
            FocalNetBackbone,
        )
        if is_torch_available()
        else ()
    )
    lowerCAmelCase__      				=							(
        {"feature-extraction": FocalNetModel, "image-classification": FocalNetForImageClassification}
        if is_torch_available()
        else {}
    )
    lowerCAmelCase__      				=							False

    lowerCAmelCase__      				=							False
    lowerCAmelCase__      				=							False
    lowerCAmelCase__      				=							False
    lowerCAmelCase__      				=							False

    def    a__				(					self     :     Dict     )							->							List[Any]:

     '''simple docstring'''




     lowerCamelCase__  =     FocalNetModelTester(self     )
     lowerCamelCase__  =     ConfigTester(self				,     config_class=__lowerCamelCase				,     embed_dim=37				,     has_text_modality=__lowerCamelCase     )

    def    a__				(					self     :     Optional[int]     )							->							Union[str, Any]:

     '''simple docstring'''




     self.create_and_test_config_common_properties()
     self.config_tester.create_and_test_config_to_json_string()
     self.config_tester.create_and_test_config_to_json_file()
     self.config_tester.create_and_test_config_from_and_save_pretrained()
     self.config_tester.create_and_test_config_with_num_labels()
     self.config_tester.check_config_can_be_init_without_params()
     self.config_tester.check_config_arguments_init()

    def    a__				(					self     :     Optional[Any]     )							->							List[Any]:

     '''simple docstring'''




     return

    def    a__				(					self     :     Union[str, Any]     )							->							Optional[int]:

     '''simple docstring'''




     lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__lowerCamelCase     )

    def    a__				(					self     :     Any     )							->							Tuple:

     '''simple docstring'''




     lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_backbone(*__lowerCamelCase     )

    def    a__				(					self     :     Any     )							->							int:

     '''simple docstring'''




     lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_image_modeling(*__lowerCamelCase     )

    def    a__				(					self     :     int     )							->							List[str]:

     '''simple docstring'''




     lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*__lowerCamelCase     )

    @unittest.skip(reason="FocalNet does not use inputs_embeds"     )
    def    a__				(					self     :     Union[str, Any]     )							->							int:

     '''simple docstring'''




     pass

    @unittest.skip(reason="FocalNet does not use feedforward chunking"     )
    def    a__				(					self     :     Dict     )							->							Optional[int]:

     '''simple docstring'''




     pass

    def    a__				(					self     :     Any     )							->							str:

     '''simple docstring'''




     lowerCamelCase__  ,	lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes[:-1]:
      lowerCamelCase__  =     model_class(__lowerCamelCase     )
      self.assertIsInstance(model.get_input_embeddings()				,     (nn.Module)     )
      lowerCamelCase__  =     model.get_output_embeddings()
      self.assertTrue(x is None or isinstance(__lowerCamelCase				,     nn.Linear     )     )

    def    a__				(					self     :     Union[str, Any]     )							->							int:

     '''simple docstring'''




     lowerCamelCase__  ,	lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes[:-1]:
      lowerCamelCase__  =     model_class(__lowerCamelCase     )
      lowerCamelCase__  =     inspect.signature(model.forward     )
      # signature.parameters is an OrderedDict => so arg_names order is deterministic
      lowerCamelCase__  =     [*signature.parameters.keys()]

      lowerCamelCase__  =     ["pixel_values"]
      self.assertListEqual(arg_names[:1]				,     __lowerCamelCase     )

    def    a__				(					self     :     str				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     List[str]				,     __lowerCamelCase     :     List[Any]     )							->							Optional[Any]:

     '''simple docstring'''




     lowerCamelCase__  =     model_class(__lowerCamelCase     )
     model.to(__lowerCamelCase     )
     model.eval()

     with torch.no_grad():
      lowerCamelCase__  =     model(**self._prepare_for_class(__lowerCamelCase				,     __lowerCamelCase     )     )

     lowerCamelCase__  =     outputs.hidden_states

     lowerCamelCase__  =     getattr(
         self.model_tester				,     "expected_num_hidden_layers"				,     len(self.model_tester.depths     ) + 1     )
     self.assertEqual(len(__lowerCamelCase     )				,     __lowerCamelCase     )

     # FocalNet has a different seq_length
     lowerCamelCase__  =     (
         config.patch_size
         if isinstance(config.patch_size				,     collections.abc.Iterable     )
         else (config.patch_size, config.patch_size)
     )

     lowerCamelCase__  =     (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

     self.assertListEqual(
         list(hidden_states[0].shape[-2:]     )				,     [num_patches, self.model_tester.embed_dim]				,     )

     lowerCamelCase__  =     outputs.reshaped_hidden_states
     self.assertEqual(len(__lowerCamelCase     )				,     __lowerCamelCase     )

     lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  =     reshaped_hidden_states[0].shape
     lowerCamelCase__  =     (
         reshaped_hidden_states[0].view(__lowerCamelCase				,     __lowerCamelCase				,     height * width     ).permute(0				,     2				,     1     )
     )
     self.assertListEqual(
         list(reshaped_hidden_states.shape[-2:]     )				,     [num_patches, self.model_tester.embed_dim]				,     )

    def    a__				(					self     :     Any     )							->							Union[str, Any]:

     '''simple docstring'''




     lowerCamelCase__  ,	lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs_for_common()

     lowerCamelCase__  =     (
         self.model_tester.image_size
         if isinstance(self.model_tester.image_size				,     collections.abc.Iterable     )
         else (self.model_tester.image_size, self.model_tester.image_size)
     )

     for model_class in self.all_model_classes[:-1]:
      lowerCamelCase__  =     True
      self.check_hidden_states_output(__lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase     )

      # check that output_hidden_states also work using config
      del inputs_dict["output_hidden_states"]
      lowerCamelCase__  =     True

      self.check_hidden_states_output(__lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase     )

    def    a__				(					self     :     Union[str, Any]     )							->							Dict:

     '''simple docstring'''




     lowerCamelCase__  ,	lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs_for_common()
     lowerCamelCase__  =     3

     lowerCamelCase__  =     (
         self.model_tester.image_size
         if isinstance(self.model_tester.image_size				,     collections.abc.Iterable     )
         else (self.model_tester.image_size, self.model_tester.image_size)
     )
     lowerCamelCase__  =     (
         config.patch_size
         if isinstance(config.patch_size				,     collections.abc.Iterable     )
         else (config.patch_size, config.patch_size)
     )

     lowerCamelCase__  =     image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
     lowerCamelCase__  =     image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

     for model_class in self.all_model_classes[:-1]:
      lowerCamelCase__  =     True
      self.check_hidden_states_output(__lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase				,     (padded_height, padded_width)     )

      # check that output_hidden_states also work using config
      del inputs_dict["output_hidden_states"]
      lowerCamelCase__  =     True
      self.check_hidden_states_output(__lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase				,     (padded_height, padded_width)     )

    @slow
    def    a__				(					self     :     int     )							->							Optional[int]:

     '''simple docstring'''




     for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      lowerCamelCase__  =     FocalNetModel.from_pretrained(__lowerCamelCase     )
      self.assertIsNotNone(__lowerCamelCase     )




    def    a__				(					self     :     Any     )							->							Optional[int]:

     '''simple docstring'''




     lowerCamelCase__  ,	lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs_for_common()

     lowerCamelCase__  =     _config_zero_init(__lowerCamelCase     )
     for model_class in self.all_model_classes:
      lowerCamelCase__  =     model_class(config=__lowerCamelCase     )
      for name, param in model.named_parameters():
       if "embeddings" not in name and param.requires_grad:
        self.assertIn(
            ((param.data.mean() * 1E9).round() / 1E9).item()				,     [0.0, 1.0]				,     msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''				,     )
@require_vision
@require_torch
class 				lowercase       (							unittest.TestCase      ):


    '''simple docstring'''

    @cached_property
    def    a__				(					self     :     Optional[int]     )							->							Tuple:

     '''simple docstring'''




     return AutoImageProcessor.from_pretrained("microsoft/focalnet-tiny"     ) if is_vision_available() else None




    @slow
    def    a__				(					self     :     List[Any]     )							->							Any:

     '''simple docstring'''




     lowerCamelCase__  =     FocalNetForImageClassification.from_pretrained("microsoft/focalnet-tiny"     ).to(__lowerCamelCase     )
     lowerCamelCase__  =     self.default_image_processor

     lowerCamelCase__  =     Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"     )
     lowerCamelCase__  =     image_processor(images=__lowerCamelCase				,     return_tensors="pt"     ).to(__lowerCamelCase     )

     # forward pass
     with torch.no_grad():
      lowerCamelCase__  =     model(**__lowerCamelCase     )

     # verify the logits
     lowerCamelCase__  =     torch.Size((1, 1000)     )
     self.assertEqual(outputs.logits.shape				,     __lowerCamelCase     )
     lowerCamelCase__  =     torch.tensor([0.2_1_6_6, -0.4_3_6_8, 0.2_1_9_1]     ).to(__lowerCamelCase     )
     self.assertTrue(torch.allclose(outputs.logits[0, :3]				,     __lowerCamelCase				,     atol=1E-4     )     )
     self.assertTrue(outputs.logits.argmax(dim=-1     ).item()				,     281     )
@require_torch
class 				lowercase       (							_lowerCamelCase ,       unittest.TestCase      ):


    '''simple docstring'''

    lowerCAmelCase__      				=							(FocalNetBackbone,) if is_torch_available() else ()
    lowerCAmelCase__      				=							FocalNetConfig

    lowerCAmelCase__      				=							False

    def    a__				(					self     :     str     )							->							Any:

     '''simple docstring'''




     lowerCamelCase__  =     FocalNetModelTester(self     )









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


__A							:							List[Any]  							=		logging.getLogger(__name__)
@dataclass
class 				lowercase       :


    '''simple docstring'''

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






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


    '''simple docstring'''

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


    '''simple docstring'''

    lowerCAmelCase__      				=							42
    lowerCAmelCase__      				=							42
    lowerCAmelCase__      				=							"longest"
    lowerCAmelCase__      				=							None
    lowerCAmelCase__      				=							None

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

     '''simple docstring'''




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

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

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

     # sample randomly masked indices
     lowerCamelCase__  =     _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 				lowercase       (							_lowerCamelCase      ):


    '''simple docstring'''

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

     '''simple docstring'''




     super().__init__(*__lowerCamelCase				,     **__lowerCamelCase     )
     lowerCamelCase__  =     0
     lowerCamelCase__  =     max_gumbel_temp
     lowerCamelCase__  =     min_gumbel_temp
     lowerCamelCase__  =     gumbel_temp_decay




    def    a__				(					self     :     Tuple				,     __lowerCamelCase     :     nn.Module				,     __lowerCamelCase     :     Dict[str, Union[torch.Tensor, Any]]     )							->							torch.Tensor:

     '''simple docstring'''




     model.train()
     lowerCamelCase__  =     self._prepare_inputs(__lowerCamelCase     )

     if self.use_amp:
      with autocast():
       lowerCamelCase__  =     self.compute_loss(__lowerCamelCase				,     __lowerCamelCase     )
     else:
      lowerCamelCase__  =     self.compute_loss(__lowerCamelCase				,     __lowerCamelCase     )

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

     if self.use_amp:
      self.scaler.scale(__lowerCamelCase     ).backward()
     elif self.use_apex:
      with amp.scale_loss(__lowerCamelCase				,     self.optimizer     ) as scaled_loss:
       scaled_loss.backward()
     elif self.deepspeed:
      self.deepspeed.backward(__lowerCamelCase     )
     else:
      loss.backward()

     self.num_update_step += 1
     # make sure gumbel softmax temperature is decayed
     if self.args.n_gpu > 1 or self.deepspeed:
      model.module.set_gumbel_temperature(
          max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step				,     self.min_gumbel_temp     )     )
     else:
      model.set_gumbel_temperature(
          max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step				,     self.min_gumbel_temp     )     )

     return loss.detach()






def    lowerCamelCase_     (    ):
 # See all possible arguments in src/transformers/training_args.py
 # or by passing the --help flag to this script.
 # We now keep distinct sets of args, for a cleaner separation of concerns.

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

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

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

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

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

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

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

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

 lowerCamelCase__  =     WavaVecaForPreTraining(lowercase__)

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

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

