





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


lowercase__		:     List[str]			=							logging.getLogger(__name__)
@dataclass
class 							__lowerCAmelCase	:






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







def    a__							(							lowercase   :	ModelArguments,   lowercase   :	TrainingArguments )						->						int:





     """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 )],   )
     _UpperCamelCase							       =			logging.WARNING
     if model_args.verbose_logging:
          _UpperCamelCase							       =			logging.DEBUG
     elif trainer_utils.is_main_process(training_args.local_rank ):
          _UpperCamelCase							       =			logging.INFO
     logger.setLevel(lowercase )
@dataclass
class 							__lowerCAmelCase	:






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

@dataclass
class 							__lowerCAmelCase	:






       """simple docstring"""
       _snake_case      :	WavaVecaForPreTraining
       _snake_case      :	WavaVecaFeatureExtractor
       _snake_case      :	Union[bool, str]									=       "longest"
       _snake_case      :	Optional[int]									=       None
       _snake_case      :	Optional[int]									=       None





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

            '''simple docstring'''




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

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

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

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

            return batch

class 							__lowerCAmelCase	(       __magic_name__					):






       """simple docstring"""





       def __init__(				self  :						str			,					*lowerCAmelCase__  :						Tuple			,					lowerCAmelCase__  :						List[str]=1			,					lowerCAmelCase__  :						List[str]=0			,					lowerCAmelCase__  :						str=1.0			,					**lowerCAmelCase__  :						List[Any]	)       ->	Tuple:

            '''simple docstring'''




            super().__init__(*lowerCAmelCase__			,					**lowerCAmelCase__	)
            _UpperCamelCase							       =			0
            _UpperCamelCase							       =			max_gumbel_temp
            _UpperCamelCase							       =			min_gumbel_temp
            _UpperCamelCase							       =			gumbel_temp_decay







       def snake_case__		(				self  :						Optional[int]			,					lowerCAmelCase__  :						nn.Module			,					lowerCAmelCase__  :						Dict[str, Union[torch.Tensor, Any]]	)       ->	torch.Tensor:

            '''simple docstring'''




            model.train()
            _UpperCamelCase							       =			self._prepare_inputs(lowerCAmelCase__	)

            if self.use_amp:
                 with autocast():
                      _UpperCamelCase							       =			self.compute_loss(lowerCAmelCase__			,					lowerCAmelCase__	)
            else:
                 _UpperCamelCase							       =			self.compute_loss(lowerCAmelCase__			,					lowerCAmelCase__	)

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

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

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

            return loss.detach()







def    a__							(							)						->						List[Any]:





     """simple docstring"""

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

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

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

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

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

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

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

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

     _UpperCamelCase							       =			WavaVecaForPreTraining(lowercase )

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

     _UpperCamelCase							       =			WavaVecaPreTrainer(
         model=lowercase,   data_collator=lowercase,   args=lowercase,   train_dataset=vectorized_datasets['''train'''],   eval_dataset=vectorized_datasets['''validation'''],   tokenizer=lowercase,   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 shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers.testing_utils import require_vision
from transformers.utils import is_vision_available


if is_vision_available():
	from PIL import Image

	from transformers import AutoProcessor, BertTokenizer, BlipImageProcessor, BlipProcessor, PreTrainedTokenizerFast




@require_vision
class       lowercase_					(unittest.TestCase   ):





				def 						__UpperCamelCase    (				self)  -> int:
											a__      =tempfile.mkdtemp()

											a__      =BlipImageProcessor()
											a__      =BertTokenizer.from_pretrained('hf-internal-testing/tiny-random-BertModel')

											a__      =BlipProcessor(lowercase_       ,  lowercase_)

											processor.save_pretrained(self.tmpdirname)





				def 						__UpperCamelCase    (				self       ,  **lowercase_)  -> str:
											return AutoProcessor.from_pretrained(self.tmpdirname       ,  **lowercase_).tokenizer





				def 						__UpperCamelCase    (				self       ,  **lowercase_)  -> List[str]:
											return AutoProcessor.from_pretrained(self.tmpdirname       ,  **lowercase_).image_processor





				def 						__UpperCamelCase    (				self)  -> Optional[int]:
											shutil.rmtree(self.tmpdirname)





				def 						__UpperCamelCase    (				self)  -> str:
											a__      =[np.random.randint(255       ,  size=(3, 30, 400)       ,  dtype=np.uinta)]

											a__      =[Image.fromarray(np.moveaxis(lowercase_       ,  0       ,  -1)) for x in image_inputs]

											return image_inputs





				def 						__UpperCamelCase    (				self)  -> str:
											a__      =BlipProcessor(tokenizer=self.get_tokenizer()       ,  image_processor=self.get_image_processor())
											processor.save_pretrained(self.tmpdirname)

											a__      =self.get_tokenizer(bos_token='(BOS)'       ,  eos_token='(EOS)')
											a__      =self.get_image_processor(do_normalize=lowercase_       ,  padding_value=1.0)

											a__      =BlipProcessor.from_pretrained(
											    self.tmpdirname       ,  bos_token='(BOS)'       ,  eos_token='(EOS)'       ,  do_normalize=lowercase_       ,  padding_value=1.0)

											self.assertEqual(processor.tokenizer.get_vocab()       ,  tokenizer_add_kwargs.get_vocab())
											self.assertIsInstance(processor.tokenizer       ,  lowercase_)

											self.assertEqual(processor.image_processor.to_json_string()       ,  image_processor_add_kwargs.to_json_string())
											self.assertIsInstance(processor.image_processor       ,  lowercase_)





				def 						__UpperCamelCase    (				self)  -> int:
											a__      =self.get_image_processor()
											a__      =self.get_tokenizer()

											a__      =BlipProcessor(tokenizer=lowercase_       ,  image_processor=lowercase_)

											a__      =self.prepare_image_inputs()

											a__      =image_processor(lowercase_       ,  return_tensors='np')
											a__      =processor(images=lowercase_       ,  return_tensors='np')

											for key in input_feat_extract.keys():
																		self.assertAlmostEqual(input_feat_extract[key].sum()       ,  input_processor[key].sum()       ,  delta=1e-2)





				def 						__UpperCamelCase    (				self)  -> List[str]:
											a__      =self.get_image_processor()
											a__      =self.get_tokenizer()

											a__      =BlipProcessor(tokenizer=lowercase_       ,  image_processor=lowercase_)

											a__      ='lower newer'

											a__      =processor(text=lowercase_)

											a__      =tokenizer(lowercase_       ,  return_token_type_ids=lowercase_)

											for key in encoded_tok.keys():
																		self.assertListEqual(encoded_tok[key]       ,  encoded_processor[key])





				def 						__UpperCamelCase    (				self)  -> int:
											a__      =self.get_image_processor()
											a__      =self.get_tokenizer()

											a__      =BlipProcessor(tokenizer=lowercase_       ,  image_processor=lowercase_)

											a__      ='lower newer'
											a__      =self.prepare_image_inputs()

											a__      =processor(text=lowercase_       ,  images=lowercase_)

											self.assertListEqual(list(inputs.keys())       ,  ['pixel_values', 'input_ids', 'attention_mask'])

											# test if it raises when no input is passed
											with pytest.raises(lowercase_):
																		processor()





				def 						__UpperCamelCase    (				self)  -> Tuple:
											a__      =self.get_image_processor()
											a__      =self.get_tokenizer()

											a__      =BlipProcessor(tokenizer=lowercase_       ,  image_processor=lowercase_)

											a__      =[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

											a__      =processor.batch_decode(lowercase_)
											a__      =tokenizer.batch_decode(lowercase_)

											self.assertListEqual(lowercase_       ,  lowercase_)





				def 						__UpperCamelCase    (				self)  -> List[Any]:
											a__      =self.get_image_processor()
											a__      =self.get_tokenizer()

											a__      =BlipProcessor(tokenizer=lowercase_       ,  image_processor=lowercase_)

											a__      ='lower newer'
											a__      =self.prepare_image_inputs()

											a__      =processor(text=lowercase_       ,  images=lowercase_)

											# For now the processor supports only ['pixel_values', 'input_ids', 'attention_mask']
											self.assertListEqual(list(inputs.keys())       ,  ['pixel_values', 'input_ids', 'attention_mask'])

