

"""simple docstring"""


from collections import defaultdict
from typing import Optional

from ..image_utils import load_image
from ..utils import (
    add_end_docstrings,
    is_torch_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_torch_available():
  import torch

  from ..models.auto.modeling_auto import MODEL_FOR_MASK_GENERATION_MAPPING

UpperCAmelCase													=logging.get_logger(__name__)





@add_end_docstrings(SCREAMING_SNAKE_CASE    )
class lowerCamelCase__   (				SCREAMING_SNAKE_CASE    ):



    '''simple docstring'''




    def __init__( self ,**lowerCamelCase_    )    ->      str:
          super().__init__(**lowerCamelCase_    )
          requires_backends(self ,"""vision"""    )
          requires_backends(self ,"""torch"""    )

          if self.framework != "pt":
                raise ValueError(f'The {self.__class__} is only available in PyTorch.'    )

          self.check_model_type(lowerCamelCase_    )




    def UpperCamelCase__		( self ,**lowerCamelCase_    )    ->      List[Any]:
          A				  =		{}
          A				  =		{}
          A				  =		{}
          # preprocess args
          if "points_per_batch" in kwargs:
                A				  =		kwargs["""points_per_batch"""]
          if "points_per_crop" in kwargs:
                A				  =		kwargs["""points_per_crop"""]
          if "crops_n_layers" in kwargs:
                A				  =		kwargs["""crops_n_layers"""]
          if "crop_overlap_ratio" in kwargs:
                A				  =		kwargs["""crop_overlap_ratio"""]
          if "crop_n_points_downscale_factor" in kwargs:
                A				  =		kwargs["""crop_n_points_downscale_factor"""]
          # postprocess args
          if "pred_iou_thresh" in kwargs:
                A				  =		kwargs["""pred_iou_thresh"""]
          if "stability_score_offset" in kwargs:
                A				  =		kwargs["""stability_score_offset"""]
          if "mask_threshold" in kwargs:
                A				  =		kwargs["""mask_threshold"""]
          if "stability_score_thresh" in kwargs:
                A				  =		kwargs["""stability_score_thresh"""]
          if "crops_nms_thresh" in kwargs:
                A				  =		kwargs["""crops_nms_thresh"""]
          if "output_rle_mask" in kwargs:
                A				  =		kwargs["""output_rle_mask"""]
          if "output_bboxes_mask" in kwargs:
                A				  =		kwargs["""output_bboxes_mask"""]
          return preprocess_kwargs, forward_params, postprocess_kwargs




    def __call__( self ,lowerCamelCase_ ,*lowerCamelCase_ ,lowerCamelCase_=None ,lowerCamelCase_=None ,**lowerCamelCase_    )    ->      Optional[Any]:
          return super().__call__(lowerCamelCase_ ,*lowerCamelCase_ ,num_workers=lowerCamelCase_ ,batch_size=lowerCamelCase_ ,**lowerCamelCase_    )




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_=6_4 ,lowerCamelCase_ = 0 ,lowerCamelCase_ = 5_1_2 / 1_5_0_0 ,lowerCamelCase_ = 3_2 ,lowerCamelCase_ = 1 ,)    ->      List[str]:
          A				  =		load_image(lowerCamelCase_    )
          A				  =		self.image_processor.size["""longest_edge"""]
          A  ,			A  ,			A  ,			A				  =		self.image_processor.generate_crop_boxes(
              lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )
          A				  =		self.image_processor(images=lowerCamelCase_ ,return_tensors="""pt"""    )

          with self.device_placement():
                if self.framework == "pt":
                      A				  =		self.get_inference_context()
                      with inference_context():
                            A				  =		self._ensure_tensor_on_device(lowerCamelCase_ ,device=self.device    )
                            A				  =		self.model.get_image_embeddings(model_inputs.pop("""pixel_values"""    )    )
                            A				  =		image_embeddings

          A				  =		grid_points.shape[1]
          A				  =		points_per_batch if points_per_batch is not None else n_points

          if points_per_batch <= 0:
                raise ValueError(
                    """Cannot have points_per_batch<=0. Must be >=1 to returned batched outputs. """
                    """To return all points at once, set points_per_batch to None"""    )

          for i in range(0 ,lowerCamelCase_ ,lowerCamelCase_    ):
                A				  =		grid_points[:, i : i + points_per_batch, :, :]
                A				  =		input_labels[:, i : i + points_per_batch]
                A				  =		i == n_points - points_per_batch
                yield {
                    "input_points": batched_points,
                    "input_labels": labels,
                    "input_boxes": crop_boxes,
                    "is_last": is_last,
                    **model_inputs,
                }




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_=0.88 ,lowerCamelCase_=0.95 ,lowerCamelCase_=0 ,lowerCamelCase_=1 ,)    ->      Tuple:
          A				  =		model_inputs.pop("""input_boxes"""    )
          A				  =		model_inputs.pop("""is_last"""    )
          A				  =		model_inputs.pop("""original_sizes"""    ).tolist()
          A				  =		model_inputs.pop("""reshaped_input_sizes"""    ).tolist()

          A				  =		self.model(**lowerCamelCase_    )

          # post processing happens here in order to avoid CPU GPU copies of ALL the masks
          A				  =		model_outputs["""pred_masks"""]
          A				  =		self.image_processor.post_process_masks(
              lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,binarize=lowerCamelCase_    )
          A				  =		model_outputs["""iou_scores"""]
          A  ,			A  ,			A				  =		self.image_processor.filter_masks(
              masks[0] ,iou_scores[0] ,original_sizes[0] ,input_boxes[0] ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,)
          return {
              "masks": masks,
              "is_last": is_last,
              "boxes": boxes,
              "iou_scores": iou_scores,
          }




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_=False ,lowerCamelCase_=False ,lowerCamelCase_=0.7 ,)    ->      Union[str, Any]:
          A				  =		[]
          A				  =		[]
          A				  =		[]
          for model_output in model_outputs:
                all_scores.append(model_output.pop("""iou_scores"""    )    )
                all_masks.extend(model_output.pop("""masks"""    )    )
                all_boxes.append(model_output.pop("""boxes"""    )    )

          A				  =		torch.cat(lowerCamelCase_    )
          A				  =		torch.cat(lowerCamelCase_    )
          A  ,			A  ,			A  ,			A				  =		self.image_processor.post_process_for_mask_generation(
              lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )

          A				  =		defaultdict(lowerCamelCase_    )
          for output in model_outputs:
                for k, v in output.items():
                      extra[k].append(lowerCamelCase_    )

          A				  =		{}
          if output_rle_mask:
                A				  =		rle_mask

          if output_bboxes_mask:
                A				  =		bounding_boxes

          return {"masks": output_masks, "scores": iou_scores, **optional, **extra}







"""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"):
  UpperCAmelCase													=True
  from torch.cuda.amp import autocast


UpperCAmelCase													=logging.getLogger(__name__)





@dataclass
class lowerCamelCase__   :



    '''simple docstring'''




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




def   _A					(	_a	:						ModelArguments     ,   _a	:						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				)]     ,   )
      A				  =		logging.WARNING
      if model_args.verbose_logging:
            A				  =		logging.DEBUG
      elif trainer_utils.is_main_process(training_args.local_rank				):
            A				  =		logging.INFO
      logger.setLevel(_a				)





@dataclass
class lowerCamelCase__   :



    '''simple docstring'''




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





@dataclass
class lowerCamelCase__   :



    '''simple docstring'''




    _lowerCamelCase					=   42
    _lowerCamelCase					=   42
    _lowerCamelCase					=   "longest"
    _lowerCamelCase					=   None
    _lowerCamelCase					=   None




    def __call__( self ,lowerCamelCase_    )    ->      Dict[str, torch.Tensor]:
          # reformat list to dict and set to pytorch format
          A				  =		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""" ,)
          A				  =		self.model._get_feat_extract_output_lengths(batch["""input_values"""].shape[-1]    )

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

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

          # sample randomly masked indices
          A				  =		_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__   (				SCREAMING_SNAKE_CASE    ):



    '''simple docstring'''




    def __init__( self ,*lowerCamelCase_ ,lowerCamelCase_=1 ,lowerCamelCase_=0 ,lowerCamelCase_=1.0 ,**lowerCamelCase_    )    ->      Union[str, Any]:
          super().__init__(*lowerCamelCase_ ,**lowerCamelCase_    )
          A				  =		0
          A				  =		max_gumbel_temp
          A				  =		min_gumbel_temp
          A				  =		gumbel_temp_decay




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_    )    ->      torch.Tensor:
          model.train()
          A				  =		self._prepare_inputs(lowerCamelCase_    )

          if self.use_amp:
                with autocast():
                      A				  =		self.compute_loss(lowerCamelCase_ ,lowerCamelCase_    )
          else:
                A				  =		self.compute_loss(lowerCamelCase_ ,lowerCamelCase_    )

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



      """simple docstring"""






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

      A  ,			A  ,			A				  =		parser.parse_args_into_dataclasses()
      configure_logger(_a     ,   _a				)

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

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

      # load audio files into numpy arrays
      A				  =		datasets.map(
          _a     ,   num_proc=data_args.preprocessing_num_workers     ,   remove_columns=datasets["""train"""].column_names				)

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

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

      # normalize and transform to `BatchFeatures`
      A				  =		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
      A				  =		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'"""				)

      A				  =		WavaVecaForPreTraining(_a				)

      A				  =		DataCollatorForWavaVecaPretraining(model=_a     ,   feature_extractor=_a				)

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




