


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


lowerCamelCase__							:						List[Any]							     =   logging.getLogger(__name__)
@dataclass
class 		__magic_name__						:





     '''simple docstring'''



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







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase				)    ->				Dict:
  logging.basicConfig(
      format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s'''							,   datefmt='''%m/%d/%Y %H:%M:%S'''							,   handlers=[logging.StreamHandler(sys.stdout				)]							,   )
  snake_case__						= logging.WARNING
  if model_args.verbose_logging:
    snake_case__						= logging.DEBUG
  elif trainer_utils.is_main_process(training_args.local_rank				):
    snake_case__						= logging.INFO
  logger.setLevel(__lowerCAmelCase				)
@dataclass
class 		__magic_name__						:





     '''simple docstring'''



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




@dataclass
class 		__magic_name__						:





     '''simple docstring'''



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

     def __call__(					self:Optional[Any]     ,		_a:List[Dict[str, Union[List[int], torch.Tensor]]]   ):
       # reformat list to dict and set to pytorch format
       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'''     ,		)
       snake_case__						= self.model._get_feat_extract_output_lengths(batch['''input_values'''].shape[-1]   )

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

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

       # sample randomly masked indices
       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 		__magic_name__						(snake_case_				):





     '''simple docstring'''



     def __init__(					self:str     ,		*_a:List[Any]     ,		_a:Dict=1     ,		_a:List[str]=0     ,		_a:Union[str, Any]=1.0     ,		**_a:Optional[Any]   ):
       super().__init__(*_a     ,		**_a   )
       snake_case__						= 0
       snake_case__						= max_gumbel_temp
       snake_case__						= min_gumbel_temp
       snake_case__						= gumbel_temp_decay







     def 	SCREAMING_SNAKE_CASE__						(					self:Tuple     ,		_a:nn.Module     ,		_a:Dict[str, Union[torch.Tensor, Any]]   ):
       model.train()
       snake_case__						= self._prepare_inputs(_a   )

       if self.use_amp:
         with autocast():
           snake_case__						= self.compute_loss(_a     ,		_a   )
       else:
         snake_case__						= self.compute_loss(_a     ,		_a   )

       if self.args.n_gpu > 1 or self.deepspeed:
         if model.module.config.ctc_loss_reduction == "mean":
           snake_case__						= loss.mean()
         elif model.module.config.ctc_loss_reduction == "sum":
           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:
         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 SCREAMING_SNAKE_CASE	(			)    ->				Union[str, Any]:
  # 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.

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

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

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

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

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

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

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

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

  snake_case__						= WavaVecaForPreTraining(__lowerCAmelCase				)

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

  snake_case__						= WavaVecaPreTrainer(
      model=__lowerCAmelCase							,   data_collator=__lowerCAmelCase							,   args=__lowerCAmelCase							,   train_dataset=vectorized_datasets['''train''']							,   eval_dataset=vectorized_datasets['''validation''']							,   tokenizer=__lowerCAmelCase							,   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()





'''simple docstring'''
from pathlib import Path
from typing import List

from transformers import is_torch_available, is_vision_available
from transformers.testing_utils import get_tests_dir, is_tool_test
from transformers.tools.agent_types import AGENT_TYPE_MAPPING, AgentAudio, AgentImage, AgentText


if is_torch_available():
					import torch

if is_vision_available():
					from PIL import Image


lowerCAmelCase   : List[Any]             =				["""text""", """image""", """audio"""]




def 	_A  (       A							)  ->				Dict:
	lowercase						:		str   					=		[]

	for input_type in input_types:
		if input_type == "text":
			inputs.append("Text input"							)
		elif input_type == "image":
			inputs.append(
			    Image.open(Path(get_tests_dir("fixtures/tests_samples/COCO"							)							) / "000000039769.png"							).resize((5_1_2, 5_1_2)							)							)
		elif input_type == "audio":
			inputs.append(torch.ones(3_0_0_0							)							)
		elif isinstance(A							,A							):
			inputs.append(create_inputs(A							)							)
		else:
			raise ValueError(F'''Invalid type requested: {input_type}'''							)

	return inputs




def 	_A  (       A							)  ->				str:
	lowercase						:		Tuple   					=		[]

	for output in outputs:
		if isinstance(A							,(str, AgentText)							):
			output_types.append("text"							)
		elif isinstance(A							,(Image.Image, AgentImage)							):
			output_types.append("image"							)
		elif isinstance(A							,(torch.Tensor, AgentAudio)							):
			output_types.append("audio"							)
		else:
			raise ValueError(F'''Invalid output: {output}'''							)

	return output_types






@is_tool_test
class        _UpperCamelCase	:


				'''simple docstring'''


				def 	a__		(     self							)    ->  Optional[Any]:
					self.assertTrue(hasattr(self.tool							,      "inputs"							)							)
					self.assertTrue(hasattr(self.tool							,      "outputs"							)							)

					lowercase						:		Optional[Any]   					=		self.tool.inputs
					for _input in inputs:
						if isinstance(_input							,      a_							):
							for __input in _input:
								self.assertTrue(__input in authorized_types							)
						else:
							self.assertTrue(_input in authorized_types							)

					lowercase						:		Any   					=		self.tool.outputs
					for _output in outputs:
						self.assertTrue(_output in authorized_types							)

				def 	a__		(     self							)    ->  Any:
					lowercase						:		Any   					=		create_inputs(self.tool.inputs							)
					lowercase						:		Tuple   					=		self.tool(*a_							)

					# There is a single output
					if len(self.tool.outputs							) == 1:
						lowercase						:		Any   					=		[outputs]

					self.assertListEqual(output_types(a_							)							,      self.tool.outputs							)

				def 	a__		(     self							)    ->  List[str]:
					self.assertTrue(hasattr(self.tool							,      "description"							)							)
					self.assertTrue(hasattr(self.tool							,      "default_checkpoint"							)							)
					self.assertTrue(self.tool.description.startswith("This is a tool that"							)							)

				def 	a__		(     self							)    ->  int:
					lowercase						:		str   					=		create_inputs(self.tool.inputs							)
					lowercase						:		str   					=		self.tool(*a_							)

					if not isinstance(a_							,      a_							):
						lowercase						:		Union[str, Any]   					=		[outputs]

					self.assertEqual(len(a_							)							,      len(self.tool.outputs							)							)

					for output, output_type in zip(a_							,      self.tool.outputs							):
						lowercase						:		List[str]   					=		AGENT_TYPE_MAPPING[output_type]
						self.assertTrue(isinstance(a_							,      a_							)							)


				def 	a__		(     self							)    ->  Optional[int]:
					lowercase						:		int   					=		create_inputs(self.tool.inputs							)

					lowercase						:		str   					=		[]

					for _input, input_type in zip(a_							,      self.tool.inputs							):
						if isinstance(a_							,      a_							):
							_inputs.append([AGENT_TYPE_MAPPING[_input_type](_input							) for _input_type in input_type]							)
						else:
							_inputs.append(AGENT_TYPE_MAPPING[input_type](_input							)							)

        # Should not raise an error
					lowercase						:		Optional[int]   					=		self.tool(*a_							)

					if not isinstance(a_							,      a_							):
						lowercase						:		str   					=		[outputs]

					self.assertEqual(len(a_							)							,      len(self.tool.outputs							)							)


