



import json
import os
import pickle
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
			import faiss

@require_faiss
class        UpperCAmelCase_  (       UpperCamelCase_	):

						'''simple docstring'''




						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= tempfile.mkdtemp()
													__SCREAMING_SNAKE_CASE									= 8

													# DPR tok
													__SCREAMING_SNAKE_CASE									= [
													    '[UNK]',
													    '[CLS]',
													    '[SEP]',
													    '[PAD]',
													    '[MASK]',
													    'want',
													    '##want',
													    '##ed',
													    'wa',
													    'un',
													    'runn',
													    '##ing',
													    ',',
													    'low',
													    'lowest',
													]
													__SCREAMING_SNAKE_CASE									= os.path.join(self.tmpdirname				,							'dpr_tokenizer'		)
													os.makedirs(_A				,							exist_ok=_A		)
													__SCREAMING_SNAKE_CASE									= os.path.join(_A				,							DPR_VOCAB_FILES_NAMES['vocab_file']		)
													with open(self.vocab_file				,							'w'				,							encoding='utf-8'		) as vocab_writer:
																				vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]		)		)

													# BART tok
													__SCREAMING_SNAKE_CASE									= [
													    'l',
													    'o',
													    'w',
													    'e',
													    'r',
													    's',
													    't',
													    'i',
													    'd',
													    'n',
													    '\u0120',
													    '\u0120l',
													    '\u0120n',
													    '\u0120lo',
													    '\u0120low',
													    'er',
													    '\u0120lowest',
													    '\u0120newer',
													    '\u0120wider',
													    '<unk>',
													]
													__SCREAMING_SNAKE_CASE									= dict(zip(_A				,							range(len(_A		)		)		)		)
													__SCREAMING_SNAKE_CASE									= ['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
													__SCREAMING_SNAKE_CASE									= {'unk_token': '<unk>'}

													__SCREAMING_SNAKE_CASE									= os.path.join(self.tmpdirname				,							'bart_tokenizer'		)
													os.makedirs(_A				,							exist_ok=_A		)
													__SCREAMING_SNAKE_CASE									= os.path.join(_A				,							BART_VOCAB_FILES_NAMES['vocab_file']		)
													__SCREAMING_SNAKE_CASE									= os.path.join(_A				,							BART_VOCAB_FILES_NAMES['merges_file']		)
													with open(self.vocab_file				,							'w'				,							encoding='utf-8'		) as fp:
																				fp.write(json.dumps(_A		) + '\n'		)
													with open(self.merges_file				,							'w'				,							encoding='utf-8'		) as fp:
																				fp.write('\n'.join(_A		)		)






						def 	_A     (   self		):







													'''simple docstring'''

													return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname				,							'dpr_tokenizer'		)		)






						def 	_A     (   self		):







													'''simple docstring'''

													return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname				,							'dpr_tokenizer'		)		)






						def 	_A     (   self		):







													'''simple docstring'''

													return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname				,							'bart_tokenizer'		)		)






						def 	_A     (   self		):







													'''simple docstring'''

													shutil.rmtree(self.tmpdirname		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= Dataset.from_dict(
													    {
													        'id': ['0', '1'],
													        'text': ['foo', 'bar'],
													        'title': ['Foo', 'Bar'],
													        'embeddings': [np.ones(self.retrieval_vector_size		), 2 * np.ones(self.retrieval_vector_size		)],
													    }		)
													dataset.add_faiss_index('embeddings'				,							string_factory='Flat'				,							metric_type=faiss.METRIC_INNER_PRODUCT		)
													return dataset






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.get_dummy_dataset()
													__SCREAMING_SNAKE_CASE									= RagConfig(
													    retrieval_vector_size=self.retrieval_vector_size				,							question_encoder=DPRConfig().to_dict()				,							generator=BartConfig().to_dict()				,							)
													with patch('transformers.models.rag.retrieval_rag.load_dataset'		) as mock_load_dataset:
																				__SCREAMING_SNAKE_CASE									= dataset
																				__SCREAMING_SNAKE_CASE									= RagRetriever(
																				    _A				,							question_encoder_tokenizer=self.get_dpr_tokenizer()				,							generator_tokenizer=self.get_bart_tokenizer()				,							)
													return retriever






						def 	_A     (   self				,							_A		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.get_dummy_dataset()
													__SCREAMING_SNAKE_CASE									= RagConfig(
													    retrieval_vector_size=self.retrieval_vector_size				,							question_encoder=DPRConfig().to_dict()				,							generator=BartConfig().to_dict()				,							index_name='custom'				,							)
													if from_disk:
																				__SCREAMING_SNAKE_CASE									= os.path.join(self.tmpdirname				,							'dataset'		)
																				__SCREAMING_SNAKE_CASE									= os.path.join(self.tmpdirname				,							'index.faiss'		)
																				dataset.get_index('embeddings'		).save(os.path.join(self.tmpdirname				,							'index.faiss'		)		)
																				dataset.drop_index('embeddings'		)
																				dataset.save_to_disk(os.path.join(self.tmpdirname				,							'dataset'		)		)
																				del dataset
																				__SCREAMING_SNAKE_CASE									= RagRetriever(
																				    _A				,							question_encoder_tokenizer=self.get_dpr_tokenizer()				,							generator_tokenizer=self.get_bart_tokenizer()				,							)
													else:
																				__SCREAMING_SNAKE_CASE									= RagRetriever(
																				    _A				,							question_encoder_tokenizer=self.get_dpr_tokenizer()				,							generator_tokenizer=self.get_bart_tokenizer()				,							index=CustomHFIndex(config.retrieval_vector_size				,							_A		)				,							)
													return retriever






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= Dataset.from_dict(
													    {
													        'id': ['0', '1'],
													        'text': ['foo', 'bar'],
													        'title': ['Foo', 'Bar'],
													        'embeddings': [np.ones(self.retrieval_vector_size + 1		), 2 * np.ones(self.retrieval_vector_size + 1		)],
													    }		)
													dataset.add_faiss_index('embeddings'				,							string_factory='Flat'				,							metric_type=faiss.METRIC_INNER_PRODUCT		)

													__SCREAMING_SNAKE_CASE									= os.path.join(self.tmpdirname				,							'hf_bert_base.hnswSQ8_correct_phi_128.c_index'		)
													dataset.save_faiss_index('embeddings'				,							index_file_name + '.index.dpr'		)
													pickle.dump(dataset['id']				,							open(index_file_name + '.index_meta.dpr'				,							'wb'		)		)

													__SCREAMING_SNAKE_CASE									= os.path.join(self.tmpdirname				,							'psgs_w100.tsv.pkl'		)
													__SCREAMING_SNAKE_CASE									= {sample['id']: [sample['text'], sample['title']] for sample in dataset}
													pickle.dump(_A				,							open(_A				,							'wb'		)		)

													__SCREAMING_SNAKE_CASE									= RagConfig(
													    retrieval_vector_size=self.retrieval_vector_size				,							question_encoder=DPRConfig().to_dict()				,							generator=BartConfig().to_dict()				,							index_name='legacy'				,							index_path=self.tmpdirname				,							)
													__SCREAMING_SNAKE_CASE									= RagRetriever(
													    _A				,							question_encoder_tokenizer=self.get_dpr_tokenizer()				,							generator_tokenizer=self.get_bart_tokenizer()		)
													return retriever






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= 1
													__SCREAMING_SNAKE_CASE									= self.get_dummy_canonical_hf_index_retriever()
													__SCREAMING_SNAKE_CASE									= np.array(
													    [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]				,							dtype=np.floataa		)
													__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= retriever.retrieve(_A				,							n_docs=_A		)
													self.assertEqual(retrieved_doc_embeds.shape				,							(2, n_docs, self.retrieval_vector_size)		)
													self.assertEqual(len(_A		)				,							2		)
													self.assertEqual(sorted(doc_dicts[0]		)				,							['embeddings', 'id', 'text', 'title']		)
													self.assertEqual(len(doc_dicts[0]['id']		)				,							_A		)
													self.assertEqual(doc_dicts[0]['id'][0]				,							'1'		)  # max inner product is reached with second doc
													self.assertEqual(doc_dicts[1]['id'][0]				,							'0'		)  # max inner product is reached with first doc
													self.assertListEqual(doc_ids.tolist()				,							[[1], [0]]		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.get_dummy_canonical_hf_index_retriever()
													with tempfile.TemporaryDirectory() as tmp_dirname:
																				with patch('transformers.models.rag.retrieval_rag.load_dataset'		) as mock_load_dataset:
																											__SCREAMING_SNAKE_CASE									= self.get_dummy_dataset()
																											retriever.save_pretrained(_A		)
																											__SCREAMING_SNAKE_CASE									= RagRetriever.from_pretrained(_A		)
																											self.assertIsInstance(_A				,							_A		)
																											__SCREAMING_SNAKE_CASE									= np.array(
																											    [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]				,							dtype=np.floataa		)
																											__SCREAMING_SNAKE_CASE									= retriever.retrieve(_A				,							n_docs=1		)
																											self.assertTrue(out is not None		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= 1
													__SCREAMING_SNAKE_CASE									= self.get_dummy_custom_hf_index_retriever(from_disk=_A		)
													__SCREAMING_SNAKE_CASE									= np.array(
													    [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]				,							dtype=np.floataa		)
													__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= retriever.retrieve(_A				,							n_docs=_A		)
													self.assertEqual(retrieved_doc_embeds.shape				,							(2, n_docs, self.retrieval_vector_size)		)
													self.assertEqual(len(_A		)				,							2		)
													self.assertEqual(sorted(doc_dicts[0]		)				,							['embeddings', 'id', 'text', 'title']		)
													self.assertEqual(len(doc_dicts[0]['id']		)				,							_A		)
													self.assertEqual(doc_dicts[0]['id'][0]				,							'1'		)  # max inner product is reached with second doc
													self.assertEqual(doc_dicts[1]['id'][0]				,							'0'		)  # max inner product is reached with first doc
													self.assertListEqual(doc_ids.tolist()				,							[[1], [0]]		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.get_dummy_custom_hf_index_retriever(from_disk=_A		)
													with tempfile.TemporaryDirectory() as tmp_dirname:
																				retriever.save_pretrained(_A		)
																				__SCREAMING_SNAKE_CASE									= RagRetriever.from_pretrained(_A		)
																				self.assertIsInstance(_A				,							_A		)
																				__SCREAMING_SNAKE_CASE									= np.array(
																				    [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]				,							dtype=np.floataa		)
																				__SCREAMING_SNAKE_CASE									= retriever.retrieve(_A				,							n_docs=1		)
																				self.assertTrue(out is not None		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= 1
													__SCREAMING_SNAKE_CASE									= self.get_dummy_custom_hf_index_retriever(from_disk=_A		)
													__SCREAMING_SNAKE_CASE									= np.array(
													    [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]				,							dtype=np.floataa		)
													__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= retriever.retrieve(_A				,							n_docs=_A		)
													self.assertEqual(retrieved_doc_embeds.shape				,							(2, n_docs, self.retrieval_vector_size)		)
													self.assertEqual(len(_A		)				,							2		)
													self.assertEqual(sorted(doc_dicts[0]		)				,							['embeddings', 'id', 'text', 'title']		)
													self.assertEqual(len(doc_dicts[0]['id']		)				,							_A		)
													self.assertEqual(doc_dicts[0]['id'][0]				,							'1'		)  # max inner product is reached with second doc
													self.assertEqual(doc_dicts[1]['id'][0]				,							'0'		)  # max inner product is reached with first doc
													self.assertListEqual(doc_ids.tolist()				,							[[1], [0]]		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.get_dummy_custom_hf_index_retriever(from_disk=_A		)
													with tempfile.TemporaryDirectory() as tmp_dirname:
																				retriever.save_pretrained(_A		)
																				__SCREAMING_SNAKE_CASE									= RagRetriever.from_pretrained(_A		)
																				self.assertIsInstance(_A				,							_A		)
																				__SCREAMING_SNAKE_CASE									= np.array(
																				    [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]				,							dtype=np.floataa		)
																				__SCREAMING_SNAKE_CASE									= retriever.retrieve(_A				,							n_docs=1		)
																				self.assertTrue(out is not None		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= 1
													__SCREAMING_SNAKE_CASE									= self.get_dummy_legacy_index_retriever()
													__SCREAMING_SNAKE_CASE									= np.array(
													    [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]				,							dtype=np.floataa		)
													__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= retriever.retrieve(_A				,							n_docs=_A		)
													self.assertEqual(retrieved_doc_embeds.shape				,							(2, n_docs, self.retrieval_vector_size)		)
													self.assertEqual(len(_A		)				,							2		)
													self.assertEqual(sorted(doc_dicts[0]		)				,							['text', 'title']		)
													self.assertEqual(len(doc_dicts[0]['text']		)				,							_A		)
													self.assertEqual(doc_dicts[0]['text'][0]				,							'bar'		)  # max inner product is reached with second doc
													self.assertEqual(doc_dicts[1]['text'][0]				,							'foo'		)  # max inner product is reached with first doc
													self.assertListEqual(doc_ids.tolist()				,							[[1], [0]]		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.get_dummy_legacy_index_retriever()
													with tempfile.TemporaryDirectory() as tmp_dirname:
																				retriever.save_pretrained(_A		)
																				__SCREAMING_SNAKE_CASE									= RagRetriever.from_pretrained(_A		)
																				self.assertIsInstance(_A				,							_A		)
																				__SCREAMING_SNAKE_CASE									= np.array(
																				    [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]				,							dtype=np.floataa		)
																				__SCREAMING_SNAKE_CASE									= retriever.retrieve(_A				,							n_docs=1		)
																				self.assertTrue(out is not None		)






						@require_torch
						@require_tokenizers
						@require_sentencepiece
						def 	_A     (   self		):







													'''simple docstring'''

													import torch

													__SCREAMING_SNAKE_CASE									= 1
													__SCREAMING_SNAKE_CASE									= self.get_dummy_canonical_hf_index_retriever()
													__SCREAMING_SNAKE_CASE									= [[5, 7], [10, 11]]
													__SCREAMING_SNAKE_CASE									= np.array(
													    [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]				,							dtype=np.floataa		)
													__SCREAMING_SNAKE_CASE									= retriever(_A				,							_A				,							prefix=retriever.config.generator.prefix				,							n_docs=_A		)
													__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= (
													    out['context_input_ids'],
													    out['context_attention_mask'],
													    out['retrieved_doc_embeds'],
													)
													self.assertEqual(retrieved_doc_embeds.shape				,							(2, n_docs, self.retrieval_vector_size)		)
													self.assertIsInstance(_A				,							_A		)
													self.assertIsInstance(_A				,							_A		)
													self.assertIsInstance(_A				,							np.ndarray		)

													__SCREAMING_SNAKE_CASE									= retriever(
													    _A				,							_A				,							prefix=retriever.config.generator.prefix				,							n_docs=_A				,							return_tensors='pt'				,							)
													__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= (  # noqa: F841
													    out['context_input_ids'],
													    out['context_attention_mask'],
													    out['retrieved_doc_embeds'],
													    out['doc_ids'],
													)
													self.assertEqual(retrieved_doc_embeds.shape				,							(2, n_docs, self.retrieval_vector_size)		)
													self.assertIsInstance(_A				,							torch.Tensor		)
													self.assertIsInstance(_A				,							torch.Tensor		)
													self.assertIsInstance(_A				,							torch.Tensor		)






						@require_torch
						@require_tokenizers
						@require_sentencepiece
						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.get_dpr_ctx_encoder_tokenizer()
													__SCREAMING_SNAKE_CASE									= 1
													__SCREAMING_SNAKE_CASE									= self.get_dummy_custom_hf_index_retriever(from_disk=_A		)
													retriever.set_ctx_encoder_tokenizer(_A		)

													__SCREAMING_SNAKE_CASE									= [[5, 7], [10, 11]]
													__SCREAMING_SNAKE_CASE									= np.array(
													    [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]				,							dtype=np.floataa		)
													__SCREAMING_SNAKE_CASE									= retriever(_A				,							_A				,							prefix=retriever.config.generator.prefix				,							n_docs=_A		)

													self.assertEqual(
													    len(_A		)				,							6		)  # check whether the retriever output consist of 6 attributes including tokenized docs
													self.assertEqual(
													    all(k in out for k in ('tokenized_doc_ids', 'tokenized_doc_attention_mask')		)				,							_A		)  # check for doc token related keys in dictionary.











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__		:					Optional[int]							=True
			from torch.cuda.amp import autocast


lowerCAmelCase__		:					List[Any]							=logging.getLogger(__name__)

@dataclass
class        UpperCAmelCase_  :

						'''simple docstring'''




						UpperCamelCase__      :			str      =  field(
						    metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models'''}	)
						UpperCamelCase__      :			Optional[str]      =  field(
						    default=UpperCamelCase_   ,   metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co'''}   ,   )
						UpperCamelCase__      :			Optional[bool]      =  field(
						    default=UpperCamelCase_   ,   metadata={'''help''': '''Whether to freeze the feature extractor layers of the model.'''}	)
						UpperCamelCase__      :			Optional[bool]      =  field(
						    default=UpperCamelCase_   ,   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.9_9_9_9_9_5   ,   metadata={'''help''': '''Decay of gumbel temperature during training.'''}	)







def      __lowercase       (			a__							,       a__       )    ->			Dict:
							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									= logging.WARNING
							if model_args.verbose_logging:
														__SCREAMING_SNAKE_CASE									= logging.DEBUG
							elif trainer_utils.is_main_process(training_args.local_rank       ):
														__SCREAMING_SNAKE_CASE									= logging.INFO
							logger.setLevel(a__       )

@dataclass
class        UpperCAmelCase_  :

						'''simple docstring'''




						UpperCamelCase__      :			str      =  field(
						    default=UpperCamelCase_   ,   metadata={'''help''': '''The name of the dataset to use (via the datasets library).'''}	)
						UpperCamelCase__      :			Optional[str]      =  field(
						    default=UpperCamelCase_   ,   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=UpperCamelCase_   ,   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=UpperCamelCase_   ,   metadata={'''help''': '''The number of processes to use for the preprocessing.'''}   ,   )
						UpperCamelCase__      :			Optional[float]      =  field(
						    default=2_0.0   ,   metadata={'''help''': '''Filter audio files that are longer than `max_duration_in_seconds` seconds'''}	)







@dataclass
class        UpperCAmelCase_  :

						'''simple docstring'''




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






						def __call__(   self				,							_A		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.feature_extractor.pad(
													    _A				,							max_length=self.max_length				,							padding=self.padding				,							pad_to_multiple_of=self.pad_to_multiple_of				,							return_tensors='pt'				,							)
													__SCREAMING_SNAKE_CASE									= self.model._get_feat_extract_output_lengths(batch['input_values'].shape[-1]		)

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

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

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

													return batch







class        UpperCAmelCase_  (       UpperCamelCase_	):

						'''simple docstring'''




						def __init__(   self				,							*_A				,							_A=1				,							_A=0				,							_A=1.0				,							**_A		):







													'''simple docstring'''

													super().__init__(*_A				,							**_A		)
													__SCREAMING_SNAKE_CASE									= 0
													__SCREAMING_SNAKE_CASE									= max_gumbel_temp
													__SCREAMING_SNAKE_CASE									= min_gumbel_temp
													__SCREAMING_SNAKE_CASE									= gumbel_temp_decay






						def 	_A     (   self				,							_A				,							_A		):







													'''simple docstring'''

													model.train()
													__SCREAMING_SNAKE_CASE									= self._prepare_inputs(_A		)

													if self.use_amp:
																				with autocast():
																											__SCREAMING_SNAKE_CASE									= self.compute_loss(_A				,							_A		)
													else:
																				__SCREAMING_SNAKE_CASE									= self.compute_loss(_A				,							_A		)

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

													if self.use_amp:
																				self.scaler.scale(_A		).backward()
													elif self.use_apex:
																				with amp.scale_loss(_A				,							self.optimizer		) as scaled_loss:
																											scaled_loss.backward()
													elif self.deepspeed:
																				self.deepspeed.backward(_A		)
													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       (			)    ->			Tuple:
							# 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									= HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)       )

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

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

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

							# load audio files into numpy arrays
							__SCREAMING_SNAKE_CASE									= 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									= vectorized_datasets.filter(
							    lambda a__       : len(data['speech']       ) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate       )       )

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

							# normalize and transform to `BatchFeatures`
							__SCREAMING_SNAKE_CASE									= 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									= 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									= WavaVecaForPreTraining(a__       )

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

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






