




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







def 			_a   (		SCREAMING_SNAKE_CASE_							:       ModelArguments   ,   SCREAMING_SNAKE_CASE_							:       TrainingArguments  ):

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






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






						def __call__(     self	,    _A    ):




											"""simple docstring"""



											__lowerCAmelCase            = 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"	,    )
											__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=_A	,    min_masks=2	,    )

											return batch






class     a__					(      snake_case__   ):






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




											"""simple docstring"""



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






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A    ):




											"""simple docstring"""



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

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

											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(_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 			_a   (		):

					'''simple docstring'''





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

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

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

					def prepare_dataset(SCREAMING_SNAKE_CASE_							:       Optional[Any]  ):
										# 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(
					    _UpperCamelCase   ,   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 SCREAMING_SNAKE_CASE_  : len(data["speech"]  ) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate  )  )

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

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

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

					__lowerCAmelCase            = WavaVecaPreTrainer(
					    model=_UpperCamelCase   ,   data_collator=_UpperCamelCase   ,   args=_UpperCamelCase   ,   train_dataset=vectorized_datasets["train"]   ,   eval_dataset=vectorized_datasets["validation"]   ,   tokenizer=_UpperCamelCase   ,   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 transformers import MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING, is_vision_available
from transformers.pipelines import pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
					from PIL import Image
else:
					class     a__					:






											@staticmethod
											def        __SCREAMING_SNAKE_CASE(     *_A	,    **_A    ):




																"""simple docstring"""



																pass
@is_pipeline_test
@require_torch
@require_vision
class     a__					(      unittest.TestCase   ):
						_a							:					Any      =					MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A	,    _A    ):




											"""simple docstring"""



											__lowerCAmelCase            = pipeline("visual-question-answering"	,    model="hf-internal-testing/tiny-vilt-random-vqa"    )
											__lowerCAmelCase            = [
											    {
											        "image": Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"    ),
											        "question": "How many cats are there?",
											    },
											    {
											        "image": "./tests/fixtures/tests_samples/COCO/000000039769.png",
											        "question": "How many cats are there?",
											    },
											]
											return vqa_pipeline, examples






						def        __SCREAMING_SNAKE_CASE(     self	,    _A	,    _A    ):




											"""simple docstring"""



											__lowerCAmelCase            = vqa_pipeline(_A	,    top_k=1    )
											self.assertEqual(
											    _A	,    [
											        [{"score": ANY(_A    ), "answer": ANY(_A    )}],
											        [{"score": ANY(_A    ), "answer": ANY(_A    )}],
											    ]	,    )






						@require_torch
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = pipeline("visual-question-answering"	,    model="hf-internal-testing/tiny-vilt-random-vqa"    )
											__lowerCAmelCase            = "./tests/fixtures/tests_samples/COCO/000000039769.png"
											__lowerCAmelCase            = "How many cats are there?"

											__lowerCAmelCase            = vqa_pipeline(image=_A	,    question="How many cats are there?"	,    top_k=2    )
											self.assertEqual(
											    _A	,    [{"score": ANY(_A    ), "answer": ANY(_A    )}, {"score": ANY(_A    ), "answer": ANY(_A    )}]    )

											__lowerCAmelCase            = vqa_pipeline({"image": image, "question": question}	,    top_k=2    )
											self.assertEqual(
											    _A	,    [{"score": ANY(_A    ), "answer": ANY(_A    )}, {"score": ANY(_A    ), "answer": ANY(_A    )}]    )






						@slow
						@require_torch
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											__lowerCAmelCase            = pipeline("visual-question-answering"	,    model="dandelin/vilt-b32-finetuned-vqa"    )
											__lowerCAmelCase            = "./tests/fixtures/tests_samples/COCO/000000039769.png"
											__lowerCAmelCase            = "How many cats are there?"

											__lowerCAmelCase            = vqa_pipeline(image=_A	,    question=_A	,    top_k=2    )
											self.assertEqual(
											    nested_simplify(_A	,    decimals=4    )	,    [{"score": 0.87_99, "answer": "2"}, {"score": 0.2_96, "answer": "1"}]    )

											__lowerCAmelCase            = vqa_pipeline({"image": image, "question": question}	,    top_k=2    )
											self.assertEqual(
											    nested_simplify(_A	,    decimals=4    )	,    [{"score": 0.87_99, "answer": "2"}, {"score": 0.2_96, "answer": "1"}]    )

											__lowerCAmelCase            = vqa_pipeline(
											    [{"image": image, "question": question}, {"image": image, "question": question}]	,    top_k=2    )
											self.assertEqual(
											    nested_simplify(_A	,    decimals=4    )	,    [[{"score": 0.87_99, "answer": "2"}, {"score": 0.2_96, "answer": "1"}]] * 2	,    )






						@require_tf
						@unittest.skip("Visual question answering not implemented in TF"    )
						def        __SCREAMING_SNAKE_CASE(     self    ):




											"""simple docstring"""



											pass




