
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        A_   :




	'''simple docstring'''


	_UpperCamelCase     :     str							 =       field(
	    metadata={"""help""": """Path to pretrained model or model identifier from huggingface.co/models"""}   )
	_UpperCamelCase     :     Optional[str]							 =       field(
	    default=__lowerCamelCase      ,	metadata={"""help""": """Where do you want to store the pretrained models downloaded from huggingface.co"""}      ,	)
	_UpperCamelCase     :     Optional[bool]							 =       field(
	    default=__lowerCamelCase      ,	metadata={"""help""": """Whether to freeze the feature extractor layers of the model."""}   )
	_UpperCamelCase     :     Optional[bool]							 =       field(
	    default=__lowerCamelCase      ,	metadata={"""help""": """Whether to log verbose messages or not."""}      ,	)
	_UpperCamelCase     :     Optional[float]							 =       field(
	    default=2.0      ,	metadata={"""help""": """Maximum temperature for gumbel softmax."""}   )
	_UpperCamelCase     :     Optional[float]							 =       field(
	    default=0.5      ,	metadata={"""help""": """Minimum temperature for gumbel softmax."""}   )
	_UpperCamelCase     :     Optional[float]							 =       field(
	    default=0.999995      ,	metadata={"""help""": """Decay of gumbel temperature during training."""}   )






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







@dataclass
class        A_   :




	'''simple docstring'''


	_UpperCamelCase     :     str							 =       field(
	    default=__lowerCamelCase      ,	metadata={"""help""": """The name of the dataset to use (via the datasets library)."""}   )
	_UpperCamelCase     :     Optional[str]							 =       field(
	    default=__lowerCamelCase      ,	metadata={"""help""": """The configuration name of the dataset to use (via the datasets library)."""}   )
	_UpperCamelCase     :     Optional[str]							 =       field(
	    default="""train"""      ,	metadata={
	        """help""": """The name of the training data set split to use (via the datasets library). Defaults to 'train'"""
	    }      ,	)
	_UpperCamelCase     :     Optional[str]							 =       field(
	    default="""validation"""      ,	metadata={
	        """help""": (
	            """The name of the validation data set split to use (via the datasets library). Defaults to 'validation'"""
	        )
	    }      ,	)
	_UpperCamelCase     :     Optional[str]							 =       field(
	    default="""file"""      ,	metadata={"""help""": """Column in the dataset that contains speech file path. Defaults to 'file'"""}      ,	)
	_UpperCamelCase     :     bool							 =       field(
	    default=__lowerCamelCase      ,	metadata={"""help""": """Overwrite the cached preprocessed datasets or not."""}   )
	_UpperCamelCase     :     Optional[int]							 =       field(
	    default=1      ,	metadata={
	        """help""": """The percentage of the train set used as validation set in case there's no validation split"""
	    }      ,	)
	_UpperCamelCase     :     Optional[int]							 =       field(
	    default=__lowerCamelCase      ,	metadata={"""help""": """The number of processes to use for the preprocessing."""}      ,	)
	_UpperCamelCase     :     Optional[float]							 =       field(
	    default=20.0      ,	metadata={"""help""": """Filter audio files that are longer than `max_duration_in_seconds` seconds"""}   )







@dataclass
class        A_   :




	'''simple docstring'''


	_UpperCamelCase     :     WavaVecaForPreTraining
	_UpperCamelCase     :     WavaVecaFeatureExtractor
	_UpperCamelCase     :     Union[bool, str]							 =       "longest"
	_UpperCamelCase     :     Optional[int]							 =       None
	_UpperCamelCase     :     Optional[int]							 =       None





	def __call__(       self		,    snake_case       ):
		# reformat list to dict and set to pytorch format
		lowercase       =    self.feature_extractor.pad(
		    snake_case		,    max_length=self.max_length		,    padding=self.padding		,    pad_to_multiple_of=self.pad_to_multiple_of		,    return_tensors='pt'		,    )
		lowercase       =    self.model._get_feat_extract_output_lengths(batch['input_values'].shape[-1]       )

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

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

		# sample randomly masked indices
		lowercase       =    _compute_mask_indices(
		    (batch_size, mask_indices_seq_length)		,    self.model.config.mask_time_prob		,    self.model.config.mask_time_length		,    attention_mask=snake_case		,    min_masks=2		,    )

		return batch







class        A_   (			__lowerCamelCase   ):




	'''simple docstring'''


	def __init__(       self		,    *snake_case		,    snake_case=1		,    snake_case=0		,    snake_case=1.0		,    **snake_case       ):
		super().__init__(*snake_case		,    **snake_case       )
		lowercase       =    0
		lowercase       =    max_gumbel_temp
		lowercase       =    min_gumbel_temp
		lowercase       =    gumbel_temp_decay





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case       ):
		model.train()
		lowercase       =    self._prepare_inputs(snake_case       )

		if self.use_amp:
			with autocast():
				lowercase       =    self.compute_loss(snake_case		,    snake_case       )
		else:
			lowercase       =    self.compute_loss(snake_case		,    snake_case       )

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

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

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

		return loss.detach()






def 							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.

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

	lowercase  ,					lowercase  ,					lowercase       =    parser.parse_args_into_dataclasses()
	configure_logger(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)

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

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

	# load audio files into numpy arrays
	lowercase       =    datasets.map(
	    __SCREAMING_SNAKE_CASE ,							num_proc=data_args.preprocessing_num_workers ,							remove_columns=datasets['train'].column_names							)

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

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

	# normalize and transform to `BatchFeatures`
	lowercase       =    vectorized_datasets.map(
	    __SCREAMING_SNAKE_CASE ,							batched=__SCREAMING_SNAKE_CASE ,							num_proc=data_args.preprocessing_num_workers ,							load_from_cache_file=not data_args.overwrite_cache ,							remove_columns=vectorized_datasets['train'].column_names ,							)

	# pretraining is only supported for "newer" stable layer norm architecture
	# apply_spec_augment has to be True, mask_feature_prob has to be 0.0
	lowercase       =    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\''							)

	lowercase       =    WavaVecaForPreTraining(__SCREAMING_SNAKE_CASE							)

	lowercase       =    DataCollatorForWavaVecaPretraining(model=__SCREAMING_SNAKE_CASE ,							feature_extractor=__SCREAMING_SNAKE_CASE							)

	lowercase       =    WavaVecaPreTrainer(
	    model=__SCREAMING_SNAKE_CASE ,							data_collator=__SCREAMING_SNAKE_CASE ,							args=__SCREAMING_SNAKE_CASE ,							train_dataset=vectorized_datasets['train'] ,							eval_dataset=vectorized_datasets['validation'] ,							tokenizer=__SCREAMING_SNAKE_CASE ,							max_gumbel_temp=model_args.max_gumbel_temperature ,							min_gumbel_temp=model_args.min_gumbel_temperature ,							gumbel_temp_decay=model_args.gumbel_temperature_decay ,							)
	trainer.train()


if __name__ == "__main__":
	main()



import unittest
from pathlib import Path
from tempfile import NamedTemporaryFile, TemporaryDirectory

from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline
from transformers.convert_graph_to_onnx import (
    convert,
    ensure_valid_input,
    generate_identified_filename,
    infer_shapes,
    quantize,
)
from transformers.testing_utils import require_tf, require_tokenizers, require_torch, slow







class        A_   :




	'''simple docstring'''


	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case       ):
		return None







class        A_   :




	'''simple docstring'''


	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    snake_case       ):
		return None







class        A_   (			unittest.TestCase   ):




	'''simple docstring'''


	_UpperCamelCase     :     Tuple							 =       [
	    # (model_name, model_kwargs)
	    ("""bert-base-cased""", {}),
	    ("""gpt2""", {"""use_cache""": False}),  # We don't support exporting GPT2 past keys anymore
	]





	@require_tf
	@slow
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
			self._test_export(snake_case		,    'tf'		,    12		,    **snake_case       )





	@require_torch
	@slow
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
			self._test_export(snake_case		,    'pt'		,    12		,    **snake_case       )





	@require_torch
	@slow
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		from transformers import BertModel

		lowercase       =    ['[UNK]', '[SEP]', '[CLS]', '[PAD]', '[MASK]', 'some', 'other', 'words']
		with NamedTemporaryFile(mode='w+t'       ) as vocab_file:
			vocab_file.write('\n'.join(snake_case       )       )
			vocab_file.flush()
			lowercase       =    BertTokenizerFast(vocab_file.name       )

		with TemporaryDirectory() as bert_save_dir:
			lowercase       =    BertModel(BertConfig(vocab_size=len(snake_case       )       )       )
			model.save_pretrained(snake_case       )
			self._test_export(snake_case		,    'pt'		,    12		,    snake_case       )





	@require_tf
	@slow
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
			lowercase       =    self._test_export(snake_case		,    'tf'		,    12		,    **snake_case       )
			lowercase       =    quantize(Path(snake_case       )       )

			# Ensure the actual quantized model is not bigger than the original one
			if quantized_path.stat().st_size >= Path(snake_case       ).stat().st_size:
				self.fail('Quantized model is bigger than initial ONNX model'       )





	@require_torch
	@slow
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
			lowercase       =    self._test_export(snake_case		,    'pt'		,    12		,    **snake_case       )
			lowercase       =    quantize(snake_case       )

			# Ensure the actual quantized model is not bigger than the original one
			if quantized_path.stat().st_size >= Path(snake_case       ).stat().st_size:
				self.fail('Quantized model is bigger than initial ONNX model'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    snake_case=None		,    **snake_case       ):
		try:
			# Compute path
			with TemporaryDirectory() as tempdir:
				lowercase       =    Path(snake_case       ).joinpath('model.onnx'       )

			# Remove folder if exists
			if path.parent.exists():
				path.parent.rmdir()

			# Export
			convert(snake_case		,    snake_case		,    snake_case		,    snake_case		,    snake_case		,    **snake_case       )

			return path
		except Exception as e:
			self.fail(snake_case       )





	@require_torch
	@require_tokenizers
	@slow
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		from transformers import BertModel

		lowercase       =    BertModel(BertConfig.from_pretrained('lysandre/tiny-bert-random'       )       )
		lowercase       =    BertTokenizerFast.from_pretrained('lysandre/tiny-bert-random'       )
		self._test_infer_dynamic_axis(snake_case		,    snake_case		,    'pt'       )





	@require_tf
	@require_tokenizers
	@slow
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		from transformers import TFBertModel

		lowercase       =    TFBertModel(BertConfig.from_pretrained('lysandre/tiny-bert-random'       )       )
		lowercase       =    BertTokenizerFast.from_pretrained('lysandre/tiny-bert-random'       )
		self._test_infer_dynamic_axis(snake_case		,    snake_case		,    'tf'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case       ):
		lowercase       =    FeatureExtractionPipeline(snake_case		,    snake_case       )

		lowercase       =    ['input_ids', 'token_type_ids', 'attention_mask', 'output_0', 'output_1']
		lowercase  ,					lowercase  ,					lowercase  ,					lowercase       =    infer_shapes(snake_case		,    snake_case       )

		# Assert all variables are present
		self.assertEqual(len(snake_case       )		,    len(snake_case       )       )
		self.assertTrue(all(var_name in shapes for var_name in variable_names       )       )
		self.assertSequenceEqual(variable_names[:3]		,    snake_case       )
		self.assertSequenceEqual(variable_names[3:]		,    snake_case       )

		# Assert inputs are {0: batch, 1: sequence}
		for var_name in ["input_ids", "token_type_ids", "attention_mask"]:
			self.assertDictEqual(shapes[var_name]		,    {0: 'batch', 1: 'sequence'}       )

		# Assert outputs are {0: batch, 1: sequence} and {0: batch}
		self.assertDictEqual(shapes['output_0']		,    {0: 'batch', 1: 'sequence'}       )
		self.assertDictEqual(shapes['output_1']		,    {0: 'batch'}       )





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    ['input_ids', 'attention_mask', 'token_type_ids']
		lowercase       =    {'input_ids': [1, 2, 3, 4], 'attention_mask': [0, 0, 0, 0], 'token_type_ids': [1, 1, 1, 1]}
		lowercase  ,					lowercase       =    ensure_valid_input(FuncContiguousArgs()		,    snake_case		,    snake_case       )

		# Should have exactly the same number of args (all are valid)
		self.assertEqual(len(snake_case       )		,    3       )

		# Should have exactly the same input names
		self.assertEqual(set(snake_case       )		,    set(snake_case       )       )

		# Parameter should be reordered according to their respective place in the function:
		# (input_ids, token_type_ids, attention_mask)
		self.assertEqual(snake_case		,    (tokens['input_ids'], tokens['token_type_ids'], tokens['attention_mask'])       )

		# Generated args are interleaved with another args (for instance parameter "past" in GPT2)
		lowercase  ,					lowercase       =    ensure_valid_input(FuncNonContiguousArgs()		,    snake_case		,    snake_case       )

		# Should have exactly the one arg (all before the one not provided "some_other_args")
		self.assertEqual(len(snake_case       )		,    1       )
		self.assertEqual(len(snake_case       )		,    1       )

		# Should have only "input_ids"
		self.assertEqual(inputs_args[0]		,    tokens['input_ids']       )
		self.assertEqual(ordered_input_names[0]		,    'input_ids'       )





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    generate_identified_filename(Path('/home/something/my_fake_model.onnx'       )		,    '-test'       )
		self.assertEqual('/home/something/my_fake_model-test.onnx'		,    generated.as_posix()       )

