



'''simple docstring'''



import os
from pickle import UnpicklingError
from typing import Dict, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict, unflatten_dict

import transformers

from .utils import logging


lowerCAmelCase_     :       List[Any]								=						logging.get_logger(__name__)
def 					UpperCAmelCase   (						A  :				List[Any]					,		A  :				List[Any]					,		A  :				Optional[Any]					,		A  :				Optional[int]=False				):




							try:
														import torch  # noqa: F401
							except ImportError:
														logger.error(
														    '''Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see'''
														    ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
														    ''' instructions.'''				)
														raise

							if not is_sharded:
														SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			os.path.abspath(_A				)
														logger.info(F"""Loading PyTorch weights from {pt_path}"""				)

														SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			torch.load(_A					,		map_location='''cpu'''				)
														logger.info(F"""PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()				):,} parameters."""				)

														SCREAMING_SNAKE_CASE					:      str     					=			convert_pytorch_state_dict_to_flax(_A					,		_A				)
							else:
														# model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
														SCREAMING_SNAKE_CASE					:      List[str]     					=			convert_pytorch_sharded_state_dict_to_flax(_A					,		_A				)
							return flax_state_dict
def 					UpperCAmelCase   (						A  :				Tuple[str]					,		A  :				np.ndarray					,		A  :				Dict[str, jnp.ndarray]					,		A  :				str					,		):




							def is_key_or_prefix_key_in_dict(A  :				Tuple[str]				) -> bool:
														return len(set(_A				) & {key, (model_prefix,) + key}				) > 0

							# layer norm
							SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			pt_tuple_key[:-1] + ("scale",)
							if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(_A				):
														return renamed_pt_tuple_key, pt_tensor

							# batch norm layer mean
							SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			pt_tuple_key[:-1] + ("mean",)
							if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(_A				):
														return renamed_pt_tuple_key, pt_tensor

							# batch norm layer var
							SCREAMING_SNAKE_CASE					:      int     					=			pt_tuple_key[:-1] + ("var",)
							if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(_A				):
														return renamed_pt_tuple_key, pt_tensor

							# embedding
							SCREAMING_SNAKE_CASE					:      str     					=			pt_tuple_key[:-1] + ("embedding",)
							if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(_A				):
														return renamed_pt_tuple_key, pt_tensor

							# conv layer
							SCREAMING_SNAKE_CASE					:      Tuple     					=			pt_tuple_key[:-1] + ("kernel",)
							if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(_A				):
														SCREAMING_SNAKE_CASE					:      Optional[int]     					=			pt_tensor.transpose(2					,		3					,		1					,		0				)
														return renamed_pt_tuple_key, pt_tensor

							# linear layer
							SCREAMING_SNAKE_CASE					:      Dict     					=			pt_tuple_key[:-1] + ("kernel",)
							if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(_A				):
														SCREAMING_SNAKE_CASE					:      int     					=			pt_tensor.T
														return renamed_pt_tuple_key, pt_tensor

							# old PyTorch layer norm weight
							SCREAMING_SNAKE_CASE					:      Optional[int]     					=			pt_tuple_key[:-1] + ("weight",)
							if pt_tuple_key[-1] == "gamma":
														return renamed_pt_tuple_key, pt_tensor

							# old PyTorch layer norm bias
							SCREAMING_SNAKE_CASE					:      Dict     					=			pt_tuple_key[:-1] + ("bias",)
							if pt_tuple_key[-1] == "beta":
														return renamed_pt_tuple_key, pt_tensor

							# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
							SCREAMING_SNAKE_CASE					:      Optional[int]     					=			None
							if pt_tuple_key[-3::2] == ("parametrizations", "original0"):
														SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			pt_tuple_key[-2] + "_g"
							elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):
														SCREAMING_SNAKE_CASE					:      Optional[int]     					=			pt_tuple_key[-2] + "_v"
							if name is not None:
														SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			pt_tuple_key[:-3] + (name,)
														return renamed_pt_tuple_key, pt_tensor

							return pt_tuple_key, pt_tensor
def 					UpperCAmelCase   (						A  :				Tuple					,		A  :				Optional[int]				):
							SCREAMING_SNAKE_CASE					:      List[str]     					=			{k: v.numpy() for k, v in pt_state_dict.items()}

							SCREAMING_SNAKE_CASE					:      List[Any]     					=			flax_model.base_model_prefix

							# use params dict if the model contains batch norm layers
							if "params" in flax_model.params:
														SCREAMING_SNAKE_CASE					:      List[Any]     					=			flax_model.params["params"]
							else:
														SCREAMING_SNAKE_CASE					:      List[str]     					=			flax_model.params
							SCREAMING_SNAKE_CASE					:      int     					=			flatten_dict(_A				)

							# add batch_stats keys,values to dict
							if "batch_stats" in flax_model.params:
														SCREAMING_SNAKE_CASE					:      Optional[int]     					=			flatten_dict(flax_model.params['''batch_stats''']				)
														random_flax_state_dict.update(_A				)

							SCREAMING_SNAKE_CASE					:      Any     					=			{}

							SCREAMING_SNAKE_CASE					:      List[str]     					=			(model_prefix not in flax_model_params) and (
							    model_prefix in {k.split('''.'''				)[0] for k in pt_state_dict.keys()}
							)
							SCREAMING_SNAKE_CASE					:      int     					=			(model_prefix in flax_model_params) and (
							    model_prefix not in {k.split('''.'''				)[0] for k in pt_state_dict.keys()}
							)

							# Need to change some parameters name to match Flax names
							for pt_key, pt_tensor in pt_state_dict.items():
														SCREAMING_SNAKE_CASE					:      Any     					=			tuple(pt_key.split('''.'''				)				)

														# remove base model prefix if necessary
														SCREAMING_SNAKE_CASE					:      Tuple     					=			pt_tuple_key[0] == model_prefix
														if load_model_with_head_into_base_model and has_base_model_prefix:
																					SCREAMING_SNAKE_CASE					:      List[str]     					=			pt_tuple_key[1:]

														# Correctly rename weight parameters
														SCREAMING_SNAKE_CASE					:      List[Any]     					=			rename_key_and_reshape_tensor(
														    _A					,		_A					,		_A					,		_A				)

														# add model prefix if necessary
														SCREAMING_SNAKE_CASE					:      int     					=			(model_prefix,) + flax_key in random_flax_state_dict
														if load_base_model_into_model_with_head and require_base_model_prefix:
																					SCREAMING_SNAKE_CASE					:      List[str]     					=			(model_prefix,) + flax_key

														if flax_key in random_flax_state_dict:
																					if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																												raise ValueError(
																												    F"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
																												    F"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""				)

        # add batch stats if the model contains batchnorm layers
														if "batch_stats" in flax_model.params:
																					if "mean" in flax_key[-1] or "var" in flax_key[-1]:
																												SCREAMING_SNAKE_CASE					:      str     					=			jnp.asarray(_A				)
																												continue
																					# remove num_batches_tracked key
																					if "num_batches_tracked" in flax_key[-1]:
																												flax_state_dict.pop(_A					,		_A				)
																												continue

																					# also add unexpected weight so that warning is thrown
																					SCREAMING_SNAKE_CASE					:      Optional[int]     					=			jnp.asarray(_A				)

														else:
																					# also add unexpected weight so that warning is thrown
																					SCREAMING_SNAKE_CASE					:      Tuple     					=			jnp.asarray(_A				)

							return unflatten_dict(_A				)
def 					UpperCAmelCase   (						A  :				List[Any]					,		A  :				Dict				):
							import torch

							# Load the index
							SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			{}
							for shard_file in shard_filenames:
														# load using msgpack utils
														SCREAMING_SNAKE_CASE					:      int     					=			torch.load(_A				)
														SCREAMING_SNAKE_CASE					:      Tuple     					=			{k: v.numpy() for k, v in pt_state_dict.items()}

														SCREAMING_SNAKE_CASE					:      int     					=			flax_model.base_model_prefix

														# use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict
														if "batch_stats" in flax_model.params:
																					SCREAMING_SNAKE_CASE					:      int     					=			flax_model.params["params"]

																					SCREAMING_SNAKE_CASE					:      int     					=			flatten_dict(_A				)
																					random_flax_state_dict.update(flatten_dict(flax_model.params['''batch_stats''']				)				)
														else:
																					SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			flax_model.params
																					SCREAMING_SNAKE_CASE					:      int     					=			flatten_dict(_A				)

														SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			(model_prefix not in flax_model_params) and (
														    model_prefix in {k.split('''.'''				)[0] for k in pt_state_dict.keys()}
														)
														SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			(model_prefix in flax_model_params) and (
														    model_prefix not in {k.split('''.'''				)[0] for k in pt_state_dict.keys()}
														)
														# Need to change some parameters name to match Flax names
														for pt_key, pt_tensor in pt_state_dict.items():
																					SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			tuple(pt_key.split('''.'''				)				)

																					# remove base model prefix if necessary
																					SCREAMING_SNAKE_CASE					:      Dict     					=			pt_tuple_key[0] == model_prefix
																					if load_model_with_head_into_base_model and has_base_model_prefix:
																												SCREAMING_SNAKE_CASE					:      str     					=			pt_tuple_key[1:]

																					# Correctly rename weight parameters
																					SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			rename_key_and_reshape_tensor(
																					    _A					,		_A					,		_A					,		_A				)
																					# add model prefix if necessary
																					SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			(model_prefix,) + flax_key in random_flax_state_dict
																					if load_base_model_into_model_with_head and require_base_model_prefix:
																												SCREAMING_SNAKE_CASE					:      Optional[int]     					=			(model_prefix,) + flax_key

																					if flax_key in random_flax_state_dict:
																												if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																																			raise ValueError(
																																			    F"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
																																			    F"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""				)

            # add batch stats if the model contains batchnorm layers
																					if "batch_stats" in flax_model.params:
																												if "mean" in flax_key[-1]:
																																			SCREAMING_SNAKE_CASE					:      Dict     					=			jnp.asarray(_A				)
																																			continue
																												if "var" in flax_key[-1]:
																																			SCREAMING_SNAKE_CASE					:      int     					=			jnp.asarray(_A				)
																																			continue
																												# remove num_batches_tracked key
																												if "num_batches_tracked" in flax_key[-1]:
																																			flax_state_dict.pop(_A					,		_A				)
																																			continue

																												# also add unexpected weight so that warning is thrown
																												SCREAMING_SNAKE_CASE					:      Tuple     					=			jnp.asarray(_A				)

																					else:
																												# also add unexpected weight so that warning is thrown
																												SCREAMING_SNAKE_CASE					:      List[str]     					=			jnp.asarray(_A				)
							return unflatten_dict(_A				)
def 					UpperCAmelCase   (						A  :				Union[str, Any]					,		A  :				str				):
							SCREAMING_SNAKE_CASE					:      List[str]     					=			os.path.abspath(_A				)
							logger.info(F"""Loading Flax weights from {flax_checkpoint_path}"""				)

							# import correct flax class
							SCREAMING_SNAKE_CASE					:      Tuple     					=			getattr(_A					,		'''Flax''' + model.__class__.__name__				)

							# load flax weight dict
							with open(_A					,		'''rb'''				) as state_f:
														try:
																					SCREAMING_SNAKE_CASE					:      List[Any]     					=			from_bytes(_A					,		state_f.read()				)
														except UnpicklingError:
																					raise EnvironmentError(F"""Unable to convert {flax_checkpoint_path} to Flax deserializable object. """				)

							return load_flax_weights_in_pytorch_model(_A					,		_A				)






def 					UpperCAmelCase   (						A  :				Optional[int]					,		A  :				Union[str, Any]				):




							try:
														import torch  # noqa: F401
							except ImportError:
														logger.error(
														    '''Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see'''
														    ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
														    ''' instructions.'''				)
														raise

							# check if we have bf16 weights
							SCREAMING_SNAKE_CASE					:      List[Any]     					=			flatten_dict(jax.tree_util.tree_map(lambda A				: x.dtype == jnp.bfloataa					,		_A				)				).values()
							if any(_A				):
														# convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16
														# and bf16 is not fully supported in PT yet.
														logger.warning(
														    '''Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` '''
														    '''before loading those in PyTorch model.'''				)
														SCREAMING_SNAKE_CASE					:      Any     					=			jax.tree_util.tree_map(
														    lambda A				: params.astype(np.floataa				) if params.dtype == jnp.bfloataa else params					,		_A				)

							SCREAMING_SNAKE_CASE					:      int     					=			flatten_dict(_A				)
							SCREAMING_SNAKE_CASE					:      Tuple     					=			pt_model.state_dict()

							SCREAMING_SNAKE_CASE					:      Dict     					=			(pt_model.base_model_prefix in flax_state) and (
							    pt_model.base_model_prefix not in {k.split('''.'''				)[0] for k in pt_model_dict.keys()}
							)
							SCREAMING_SNAKE_CASE					:      str     					=			(pt_model.base_model_prefix not in flax_state) and (
							    pt_model.base_model_prefix in {k.split('''.'''				)[0] for k in pt_model_dict.keys()}
							)

							# keep track of unexpected & missing keys
							SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			[]
							SCREAMING_SNAKE_CASE					:      List[str]     					=			set(pt_model_dict.keys()				)

							for flax_key_tuple, flax_tensor in flax_state_dict.items():
														SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			flax_key_tuple[0] == pt_model.base_model_prefix
														SCREAMING_SNAKE_CASE					:      List[Any]     					=			".".join((pt_model.base_model_prefix,) + flax_key_tuple				) in pt_model_dict

														# adapt flax_key to prepare for loading from/to base model only
														if load_model_with_head_into_base_model and has_base_model_prefix:
																					SCREAMING_SNAKE_CASE					:      Dict     					=			flax_key_tuple[1:]
														elif load_base_model_into_model_with_head and require_base_model_prefix:
																					SCREAMING_SNAKE_CASE					:      List[str]     					=			(pt_model.base_model_prefix,) + flax_key_tuple

														# rename flax weights to PyTorch format
														if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(_A				) not in pt_model_dict:
																					# conv layer
																					SCREAMING_SNAKE_CASE					:      Dict     					=			flax_key_tuple[:-1] + ("weight",)
																					SCREAMING_SNAKE_CASE					:      str     					=			jnp.transpose(_A					,		(3, 2, 0, 1)				)
														elif flax_key_tuple[-1] == "kernel" and ".".join(_A				) not in pt_model_dict:
																					# linear layer
																					SCREAMING_SNAKE_CASE					:      Any     					=			flax_key_tuple[:-1] + ("weight",)
																					SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			flax_tensor.T
														elif flax_key_tuple[-1] in ["scale", "embedding"]:
																					SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			flax_key_tuple[:-1] + ("weight",)

														# adding batch stats from flax batch norm to pt
														elif "mean" in flax_key_tuple[-1]:
																					SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			flax_key_tuple[:-1] + ("running_mean",)
														elif "var" in flax_key_tuple[-1]:
																					SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			flax_key_tuple[:-1] + ("running_var",)

														if "batch_stats" in flax_state:
																					SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			".".join(flax_key_tuple[1:]				)  # Remove the params/batch_stats header
														else:
																					SCREAMING_SNAKE_CASE					:      int     					=			".".join(_A				)

														# We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
														SCREAMING_SNAKE_CASE					:      List[Any]     					=			{}
														# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
														for key in pt_model_dict:
																					SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			key.split('''.'''				)
																					SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			None
																					if key_components[-3::2] == ["parametrizations", "original0"]:
																												SCREAMING_SNAKE_CASE					:      str     					=			key_components[-2] + "_g"
																					elif key_components[-3::2] == ["parametrizations", "original1"]:
																												SCREAMING_SNAKE_CASE					:      int     					=			key_components[-2] + "_v"
																					if name is not None:
																												SCREAMING_SNAKE_CASE					:      Tuple     					=			key_components[:-3] + [name]
																												SCREAMING_SNAKE_CASE					:      str     					=			".".join(_A				)
																												SCREAMING_SNAKE_CASE					:      Dict     					=			key

														if flax_key in special_pt_names:
																					SCREAMING_SNAKE_CASE					:      str     					=			special_pt_names[flax_key]

														if flax_key in pt_model_dict:
																					if flax_tensor.shape != pt_model_dict[flax_key].shape:
																												raise ValueError(
																												    F"""Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected """
																												    F"""to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}."""				)
																					else:
																												# add weight to pytorch dict
																												SCREAMING_SNAKE_CASE					:      Dict     					=			np.asarray(_A				) if not isinstance(_A					,		np.ndarray				) else flax_tensor
																												SCREAMING_SNAKE_CASE					:      Tuple     					=			torch.from_numpy(_A				)
																												# remove from missing keys
																												missing_keys.remove(_A				)
														else:
																					# weight is not expected by PyTorch model
																					unexpected_keys.append(_A				)

							pt_model.load_state_dict(_A				)

							# re-transform missing_keys to list
							SCREAMING_SNAKE_CASE					:      Optional[int]     					=			list(_A				)

							if len(_A				) > 0:
														logger.warning(
														    '''Some weights of the Flax model were not used when initializing the PyTorch model'''
														    F""" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"""
														    F""" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture"""
														    ''' (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This'''
														    F""" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect"""
														    ''' to be exactly identical (e.g. initializing a BertForSequenceClassification model from a'''
														    ''' FlaxBertForSequenceClassification model).'''				)
							else:
														logger.warning(F"""All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n"""				)
							if len(_A				) > 0:
														logger.warning(
														    F"""Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly"""
														    F""" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to"""
														    ''' use it for predictions and inference.'''				)
							else:
														logger.warning(
														    F"""All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n"""
														    '''If your task is similar to the task the model of the checkpoint was trained on, '''
														    F"""you can already use {pt_model.__class__.__name__} for predictions without further training."""				)

							return pt_model









'''simple docstring'''



import logging
import sys
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

import librosa
import torch
from datasets import DatasetDict, load_dataset
from packaging import version
from torch import nn

from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaForPreTraining,
    is_apex_available,
    trainer_utils,
)
from transformers.models.wavaveca.modeling_wavaveca import _compute_mask_indices


if is_apex_available():
							from apex import amp

if version.parse(version.parse(torch.__version__).base_version) >= version.parse('1.6'):
							lowerCAmelCase_     :       Union[str, Any]								=						True
							from torch.cuda.amp import autocast


lowerCAmelCase_     :       Union[str, Any]								=						logging.getLogger(__name__)
@dataclass
class 	lowerCamelCase_							:
						_lowerCAmelCase     :   str	     =				field(
						    metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}			)
						_lowerCAmelCase     :   Optional[str]	     =				field(
						    default=snake_case_    ,				metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'}    ,				)
						_lowerCAmelCase     :   Optional[bool]	     =				field(
						    default=snake_case_    ,				metadata={'help': 'Whether to freeze the feature extractor layers of the model.'}			)
						_lowerCAmelCase     :   Optional[bool]	     =				field(
						    default=snake_case_    ,				metadata={'help': 'Whether to log verbose messages or not.'}    ,				)
						_lowerCAmelCase     :   Optional[float]	     =				field(
						    default=2.0    ,				metadata={'help': 'Maximum temperature for gumbel softmax.'}			)
						_lowerCAmelCase     :   Optional[float]	     =				field(
						    default=0.5    ,				metadata={'help': 'Minimum temperature for gumbel softmax.'}			)
						_lowerCAmelCase     :   Optional[float]	     =				field(
						    default=0.99_9995    ,				metadata={'help': 'Decay of gumbel temperature during training.'}			)


def 					UpperCAmelCase   (						A  :				ModelArguments					,		A  :				TrainingArguments				):
							logging.basicConfig(
							    format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s'''					,		datefmt='''%m/%d/%Y %H:%M:%S'''					,		handlers=[logging.StreamHandler(sys.stdout				)]					,		)
							SCREAMING_SNAKE_CASE					:      List[Any]     					=			logging.WARNING
							if model_args.verbose_logging:
														SCREAMING_SNAKE_CASE					:      Dict     					=			logging.DEBUG
							elif trainer_utils.is_main_process(training_args.local_rank				):
														SCREAMING_SNAKE_CASE					:      Any     					=			logging.INFO
							logger.setLevel(A				)
@dataclass
class 	lowerCamelCase_							:
						_lowerCAmelCase     :   str	     =				field(
						    default=snake_case_    ,				metadata={'help': 'The name of the dataset to use (via the datasets library).'}			)
						_lowerCAmelCase     :   Optional[str]	     =				field(
						    default=snake_case_    ,				metadata={'help': 'The configuration name of the dataset to use (via the datasets library).'}			)
						_lowerCAmelCase     :   Optional[str]	     =				field(
						    default='train'    ,				metadata={
						        'help': 'The name of the training data set split to use (via the datasets library). Defaults to \'train\''
						    }    ,				)
						_lowerCAmelCase     :   Optional[str]	     =				field(
						    default='validation'    ,				metadata={
						        'help': (
						            'The name of the validation data set split to use (via the datasets library). Defaults to \'validation\''
						        )
						    }    ,				)
						_lowerCAmelCase     :   Optional[str]	     =				field(
						    default='file'    ,				metadata={'help': 'Column in the dataset that contains speech file path. Defaults to \'file\''}    ,				)
						_lowerCAmelCase     :   bool	     =				field(
						    default=snake_case_    ,				metadata={'help': 'Overwrite the cached preprocessed datasets or not.'}			)
						_lowerCAmelCase     :   Optional[int]	     =				field(
						    default=1    ,				metadata={
						        'help': 'The percentage of the train set used as validation set in case there\'s no validation split'
						    }    ,				)
						_lowerCAmelCase     :   Optional[int]	     =				field(
						    default=snake_case_    ,				metadata={'help': 'The number of processes to use for the preprocessing.'}    ,				)
						_lowerCAmelCase     :   Optional[float]	     =				field(
						    default=20.0    ,				metadata={'help': 'Filter audio files that are longer than `max_duration_in_seconds` seconds'}			)


@dataclass
class 	lowerCamelCase_							:
						_lowerCAmelCase     :   WavaVecaForPreTraining
						_lowerCAmelCase     :   WavaVecaFeatureExtractor
						_lowerCAmelCase     :   Union[bool, str]	     =				"longest"
						_lowerCAmelCase     :   Optional[int]	     =				None
						_lowerCAmelCase     :   Optional[int]	     =				None







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

													"""simple docstring"""
													# reformat list to dict and set to pytorch format
													SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			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'''				,    )
													SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			self.model._get_feat_extract_output_lengths(batch['''input_values'''].shape[-1]				)

													SCREAMING_SNAKE_CASE					:      Tuple     					=			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
																				SCREAMING_SNAKE_CASE					:      List[Any]     					=			self.model._get_feat_extract_output_lengths(batch['''attention_mask'''].sum(-1				)				).to(
																				    torch.long				)

																				SCREAMING_SNAKE_CASE					:      List[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
																				SCREAMING_SNAKE_CASE					:      Optional[int]     					=			1
																				SCREAMING_SNAKE_CASE					:      Dict     					=			attention_mask.flip([-1]				).cumsum(-1				).flip([-1]				).bool()

													# sample randomly masked indices
													SCREAMING_SNAKE_CASE					:      List[str]     					=			_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_							(					snake_case_			):







						def __init__(  self   :      str				,    *lowerCAmelCase__   :      int				,    lowerCAmelCase__   :      Any=1				,    lowerCAmelCase__   :      Dict=0				,    lowerCAmelCase__   :      Tuple=1.0				,    **lowerCAmelCase__   :      Optional[Any]				):

													"""simple docstring"""
													super().__init__(*lowerCAmelCase__				,    **lowerCAmelCase__				)
													SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			0
													SCREAMING_SNAKE_CASE					:      List[Any]     					=			max_gumbel_temp
													SCREAMING_SNAKE_CASE					:      List[Any]     					=			min_gumbel_temp
													SCREAMING_SNAKE_CASE					:      Tuple     					=			gumbel_temp_decay







						def 					__lowercase    (  self   :      str				,    lowerCAmelCase__   :      nn.Module				,    lowerCAmelCase__   :      Dict[str, Union[torch.Tensor, Any]]				):

													"""simple docstring"""
													model.train()
													SCREAMING_SNAKE_CASE					:      Optional[int]     					=			self._prepare_inputs(lowerCAmelCase__				)

													if self.use_amp:
																				with autocast():
																											SCREAMING_SNAKE_CASE					:      Any     					=			self.compute_loss(lowerCAmelCase__				,    lowerCAmelCase__				)
													else:
																				SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			self.compute_loss(lowerCAmelCase__				,    lowerCAmelCase__				)

													if self.args.n_gpu > 1 or self.deepspeed:
																				if model.module.config.ctc_loss_reduction == "mean":
																											SCREAMING_SNAKE_CASE					:      Any     					=			loss.mean()
																				elif model.module.config.ctc_loss_reduction == "sum":
																											SCREAMING_SNAKE_CASE					:      Tuple     					=			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:
																				SCREAMING_SNAKE_CASE					:      List[Any]     					=			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 					UpperCAmelCase   (						):
							# 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.

							SCREAMING_SNAKE_CASE					:      Optional[int]     					=			HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)				)

							SCREAMING_SNAKE_CASE					,SCREAMING_SNAKE_CASE					,SCREAMING_SNAKE_CASE					:      List[Any]     					=			parser.parse_args_into_dataclasses()
							configure_logger(A					,		A				)

							# Downloading and loading a dataset from the hub.
							SCREAMING_SNAKE_CASE					:      str     					=			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"
														SCREAMING_SNAKE_CASE					:      Any     					=			DatasetDict()
														SCREAMING_SNAKE_CASE					:      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					,		)
														SCREAMING_SNAKE_CASE					:      List[str]     					=			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"
														SCREAMING_SNAKE_CASE					:      str     					=			DatasetDict()
														SCREAMING_SNAKE_CASE					:      Any     					=			load_dataset(
														    data_args.dataset_name					,		data_args.dataset_config_name					,		split='''validation'''					,		cache_dir=model_args.cache_dir					,		)
														SCREAMING_SNAKE_CASE					:      Tuple     					=			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
							SCREAMING_SNAKE_CASE					:      int     					=			WavaVecaFeatureExtractor.from_pretrained(
							    model_args.model_name_or_path					,		cache_dir=model_args.cache_dir					,		do_normalize=A				)

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

							# load audio files into numpy arrays
							SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			datasets.map(
							    A					,		num_proc=data_args.preprocessing_num_workers					,		remove_columns=datasets['''train'''].column_names				)

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

							def normalize(A  :				List[Any]				):
														return feature_extractor(batch['''speech''']					,		sampling_rate=feature_extractor.sampling_rate				)

							# normalize and transform to `BatchFeatures`
							SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			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
							SCREAMING_SNAKE_CASE					:      Tuple     					=			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\''''				)

							SCREAMING_SNAKE_CASE					:      str     					=			WavaVecaForPreTraining(A				)

							SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			DataCollatorForWavaVecaPretraining(model=A					,		feature_extractor=A				)

							SCREAMING_SNAKE_CASE					:      Dict     					=			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()




