

'''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      				=     True
		from torch.cuda.amp import autocast


lowerCamelCase      				=     logging.getLogger(__name__)





@dataclass
class     snake_case_			:


							"""simple docstring"""





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


def        __lowercase    (  UpperCAmelCase__		,					UpperCAmelCase__      ):

					"""simple docstring"""






					logging.basicConfig(
					    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'		,					datefmt='%m/%d/%Y %H:%M:%S'		,					handlers=[logging.StreamHandler(sys.stdout      )]		,					)
					__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(_snake_case      )





@dataclass
class     snake_case_			:


							"""simple docstring"""





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





@dataclass
class     snake_case_			:


							"""simple docstring"""





							__UpperCAmelCase       =42
							__UpperCAmelCase       =42
							__UpperCAmelCase       ="longest"
							__UpperCAmelCase       =None
							__UpperCAmelCase       =None




							def __call__( self					,						_A							):
												# reformat list to dict and set to pytorch format
												__lowerCAmelCase					=							self.feature_extractor.pad(
												    UpperCAmelCase_					,						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=UpperCAmelCase_					,						min_masks=2					,						)

												return batch





class     snake_case_			(				__UpperCAmelCase							):


							"""simple docstring"""





							def __init__( self					,						*_A					,						_A=1					,						_A=0					,						_A=1.0					,						**_A							):
												super().__init__(*UpperCAmelCase_					,						**UpperCAmelCase_							)
												__lowerCAmelCase					=							0
												__lowerCAmelCase					=							max_gumbel_temp
												__lowerCAmelCase					=							min_gumbel_temp
												__lowerCAmelCase					=							gumbel_temp_decay




							def        A__  ( self					,						_A					,						_A							):
												model.train()
												__lowerCAmelCase					=							self._prepare_inputs(UpperCAmelCase_							)

												if self.use_amp:
																	with autocast():
																						__lowerCAmelCase					=							self.compute_loss(UpperCAmelCase_					,						UpperCAmelCase_							)
												else:
																	__lowerCAmelCase					=							self.compute_loss(UpperCAmelCase_					,						UpperCAmelCase_							)

												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(UpperCAmelCase_							).backward()
												elif self.use_apex:
																	with amp.scale_loss(UpperCAmelCase_					,						self.optimizer							) as scaled_loss:
																						scaled_loss.backward()
												elif self.deepspeed:
																	self.deepspeed.backward(UpperCAmelCase_							)
												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        __lowercase    (  ):

					"""simple docstring"""






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

					__lowerCAmelCase,	__lowerCAmelCase,	__lowerCAmelCase					=							parser.parse_args_into_dataclasses()
					configure_logger(_snake_case		,					_snake_case      )

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

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

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

					# normalize and transform to `BatchFeatures`
					__lowerCAmelCase					=							vectorized_datasets.map(
					    _snake_case		,					batched=_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
					__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(_snake_case      )

					__lowerCAmelCase					=							DataCollatorForWavaVecaPretraining(model=_snake_case		,					feature_extractor=_snake_case      )

					__lowerCAmelCase					=							WavaVecaPreTrainer(
					    model=_snake_case		,					data_collator=_snake_case		,					args=_snake_case		,					train_dataset=vectorized_datasets['train']		,					eval_dataset=vectorized_datasets['validation']		,					tokenizer=_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 argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
lowerCamelCase      				=     data_utils.TransfoXLTokenizer
lowerCamelCase      				=     data_utils.TransfoXLCorpus
lowerCamelCase      				=     data_utils
lowerCamelCase      				=     data_utils


def        __lowercase    (  UpperCAmelCase__		,					UpperCAmelCase__		,					UpperCAmelCase__		,					UpperCAmelCase__      ):

					"""simple docstring"""






					if transfo_xl_dataset_file:
										# Convert a pre-processed corpus (see original TensorFlow repo)
										with open(UpperCAmelCase__		,					'rb'      ) as fp:
															__lowerCAmelCase					=							pickle.load(UpperCAmelCase__		,					encoding='latin1'      )
										# Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
										__lowerCAmelCase					=							pytorch_dump_folder_path + '/' + VOCAB_FILES_NAMES['pretrained_vocab_file']
										print(F"""Save vocabulary to {pytorch_vocab_dump_path}"""      )
										__lowerCAmelCase					=							corpus.vocab.__dict__
										torch.save(UpperCAmelCase__		,					UpperCAmelCase__      )

										__lowerCAmelCase					=							corpus.__dict__
										corpus_dict_no_vocab.pop('vocab'		,					UpperCAmelCase__      )
										__lowerCAmelCase					=							pytorch_dump_folder_path + '/' + CORPUS_NAME
										print(F"""Save dataset to {pytorch_dataset_dump_path}"""      )
										torch.save(UpperCAmelCase__		,					UpperCAmelCase__      )

					if tf_checkpoint_path:
										# Convert a pre-trained TensorFlow model
										__lowerCAmelCase					=							os.path.abspath(UpperCAmelCase__      )
										__lowerCAmelCase					=							os.path.abspath(UpperCAmelCase__      )

										print(F"""Converting Transformer XL checkpoint from {tf_path} with config at {config_path}."""      )
										# Initialise PyTorch model
										if transfo_xl_config_file == "":
															__lowerCAmelCase					=							TransfoXLConfig()
										else:
															__lowerCAmelCase					=							TransfoXLConfig.from_json_file(UpperCAmelCase__      )
										print(F"""Building PyTorch model from configuration: {config}"""      )
										__lowerCAmelCase					=							TransfoXLLMHeadModel(UpperCAmelCase__      )

										__lowerCAmelCase					=							load_tf_weights_in_transfo_xl(UpperCAmelCase__		,					UpperCAmelCase__		,					UpperCAmelCase__      )
										# Save pytorch-model
										__lowerCAmelCase					=							os.path.join(UpperCAmelCase__		,					UpperCAmelCase__      )
										__lowerCAmelCase					=							os.path.join(UpperCAmelCase__		,					UpperCAmelCase__      )
										print(F"""Save PyTorch model to {os.path.abspath(UpperCAmelCase__      )}"""      )
										torch.save(model.state_dict()		,					UpperCAmelCase__      )
										print(F"""Save configuration file to {os.path.abspath(UpperCAmelCase__      )}"""      )
										with open(UpperCAmelCase__		,					'w'		,					encoding='utf-8'      ) as f:
															f.write(config.to_json_string()      )


if __name__ == "__main__":
		lowerCamelCase      				=     argparse.ArgumentParser()
		parser.add_argument(
		    '''--pytorch_dump_folder_path''',
		    default=None,
		    type=str,
		    required=True,
		    help='''Path to the folder to store the PyTorch model or dataset/vocab.''',
		)
		parser.add_argument(
		    '''--tf_checkpoint_path''',
		    default='''''',
		    type=str,
		    help='''An optional path to a TensorFlow checkpoint path to be converted.''',
		)
		parser.add_argument(
		    '''--transfo_xl_config_file''',
		    default='''''',
		    type=str,
		    help=(
		        '''An optional config json file corresponding to the pre-trained BERT model. \n'''
		        '''This specifies the model architecture.'''
		    ),
		)
		parser.add_argument(
		    '''--transfo_xl_dataset_file''',
		    default='''''',
		    type=str,
		    help='''An optional dataset file to be converted in a vocabulary.''',
		)
		lowerCamelCase      				=     parser.parse_args()
		convert_transfo_xl_checkpoint_to_pytorch(
		    args.tf_checkpoint_path,
		    args.transfo_xl_config_file,
		    args.pytorch_dump_folder_path,
		    args.transfo_xl_dataset_file,
		)






