


'''simple docstring'''





import logging
import os
import sys
import warnings
from dataclasses import dataclass, field
from random import randint
from typing import Optional

import datasets
import evaluate
import numpy as np
from datasets import DatasetDict, load_dataset

import transformers
from transformers import (
    AutoConfig,
    AutoFeatureExtractor,
    AutoModelForAudioClassification,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version


__lowerCamelCase												=						logging.getLogger(__name__)

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version('''4.31.0''')

require_version('''datasets>=1.14.0''', '''To fix: pip install -r examples/pytorch/audio-classification/requirements.txt''')


def        a__     (					UpperCamelCase_						:      np.ndarray,   UpperCamelCase_						:      float,   UpperCamelCase_						:      int = 16_000						):
   UpperCAmelCase__	:List[str]      =      int(round(sample_rate * max_length						)						)
   if len(UpperCamelCase_						) <= sample_length:
      return wav
   UpperCAmelCase__	:int      =      randint(0,   len(UpperCamelCase_						) - sample_length - 1						)
   return wav[random_offset : random_offset + sample_length]




@dataclass
class 				UpperCAmelCase			:
 UpperCAmelCase				=					field(default=_snake_case						, metadata={"help": "Name of a dataset from the datasets package"}						)
 UpperCAmelCase				=					field(
     default=_snake_case						, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}						)
 UpperCAmelCase				=					field(
     default=_snake_case						, metadata={"help": "A file containing the training audio paths and labels."}						)
 UpperCAmelCase				=					field(
     default=_snake_case						, metadata={"help": "A file containing the validation audio paths and labels."}						)
 UpperCAmelCase				=					field(
     default="train"						, metadata={
         "help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"
     }						, )
 UpperCAmelCase				=					field(
     default="validation"						, metadata={
         "help": (
             "The name of the training data set split to use (via the datasets library). Defaults to 'validation'"
         )
     }						, )
 UpperCAmelCase				=					field(
     default="audio"						, metadata={"help": "The name of the dataset column containing the audio data. Defaults to 'audio'"}						, )
 UpperCAmelCase				=					field(
     default="label"						, metadata={"help": "The name of the dataset column containing the labels. Defaults to 'label'"}						)
 UpperCAmelCase				=					field(
     default=_snake_case						, metadata={
         "help": (
             "For debugging purposes or quicker training, truncate the number of training examples to this "
             "value if set."
         )
     }						, )
 UpperCAmelCase				=					field(
     default=_snake_case						, metadata={
         "help": (
             "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
             "value if set."
         )
     }						, )
 UpperCAmelCase				=					field(
     default=20						, metadata={"help": "Audio clips will be randomly cut to this length during training if the value is set."}						, )




@dataclass
class 				UpperCAmelCase			:
 UpperCAmelCase				=					field(
     default="facebook/wav2vec2-base"						, metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}						, )
 UpperCAmelCase				=					field(
     default=_snake_case						, metadata={"help": "Pretrained config name or path if not the same as model_name"}						)
 UpperCAmelCase				=					field(
     default=_snake_case						, metadata={"help": "Where do you want to store the pretrained models downloaded from the Hub"}						)
 UpperCAmelCase				=					field(
     default="main"						, metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."}						, )
 UpperCAmelCase				=					field(
     default=_snake_case						, metadata={"help": "Name or path of preprocessor config."}						)
 UpperCAmelCase				=					field(
     default=_snake_case						, metadata={"help": "Whether to freeze the feature encoder layers of the model."}						)
 UpperCAmelCase				=					field(
     default=_snake_case						, metadata={"help": "Whether to generate an attention mask in the feature extractor."}						)
 UpperCAmelCase				=					field(
     default=_snake_case						, metadata={
         "help": (
             "Will use the token generated when running `huggingface-cli login` (necessary to use this script "
             "with private models)."
         )
     }						, )
 UpperCAmelCase				=					field(
     default=_snake_case						, metadata={"help": "Whether to freeze the feature extractor layers of the model."}						)
 UpperCAmelCase				=					field(
     default=_snake_case						, metadata={"help": "Will enable to load a pretrained model whose head dimensions are different."}						, )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    if not self.freeze_feature_extractor and self.freeze_feature_encoder:
       warnings.warn(
           '''The argument `--freeze_feature_extractor` is deprecated and '''
           '''will be removed in a future version. Use `--freeze_feature_encoder`'''
           '''instead. Setting `freeze_feature_encoder==True`.'''			,       __lowerCamelCase			,       )
    if self.freeze_feature_extractor and not self.freeze_feature_encoder:
       raise ValueError(
           '''The argument `--freeze_feature_extractor` is deprecated and '''
           '''should not be used in combination with `--freeze_feature_encoder`.'''
           '''Only make use of `--freeze_feature_encoder`.''' )


def        a__     (					):
   # 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.

   UpperCAmelCase__	:List[Any]      =      HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)						)
   if len(sys.argv						) == 2 and sys.argv[1].endswith('''.json'''						):
      # If we pass only one argument to the script and it's the path to a json file,
      # let's parse it to get our arguments.
      UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:Optional[int]      =      parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]						)						)
   else:
      UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:List[Any]      =      parser.parse_args_into_dataclasses()

   # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
   # information sent is the one passed as arguments along with your Python/PyTorch versions.
   send_example_telemetry('''run_audio_classification''',   UpperCamelCase_,   UpperCamelCase_						)

   # Setup logging
   logging.basicConfig(
       format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s''',   datefmt='''%m/%d/%Y %H:%M:%S''',   handlers=[logging.StreamHandler(sys.stdout						)],   )

   if training_args.should_log:
      # The default of training_args.log_level is passive, so we set log level at info here to have that default.
      transformers.utils.logging.set_verbosity_info()

   UpperCAmelCase__	:Dict      =      training_args.get_process_log_level()
   logger.setLevel(UpperCamelCase_						)
   transformers.utils.logging.set_verbosity(UpperCamelCase_						)
   transformers.utils.logging.enable_default_handler()
   transformers.utils.logging.enable_explicit_format()

   # Log on each process the small summary:
   logger.warning(
       F'''Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu} '''
       + F'''distributed training: {bool(training_args.local_rank != -1						)}, 16-bits training: {training_args.fpaa}'''						)
   logger.info(F'''Training/evaluation parameters {training_args}'''						)

   # Set seed before initializing model.
   set_seed(training_args.seed						)

   # Detecting last checkpoint.
   UpperCAmelCase__	:int      =      None
   if os.path.isdir(training_args.output_dir						) and training_args.do_train and not training_args.overwrite_output_dir:
      UpperCAmelCase__	:Dict      =      get_last_checkpoint(training_args.output_dir						)
      if last_checkpoint is None and len(os.listdir(training_args.output_dir						)						) > 0:
         raise ValueError(
             F'''Output directory ({training_args.output_dir}) already exists and is not empty. '''
             '''Use --overwrite_output_dir to train from scratch.'''						)
      elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
         logger.info(
             F'''Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change '''
             '''the `--output_dir` or add `--overwrite_output_dir` to train from scratch.'''						)

    # Initialize our dataset and prepare it for the audio classification task.
   UpperCAmelCase__	:Optional[int]      =      DatasetDict()
   UpperCAmelCase__	:Tuple      =      load_dataset(
       data_args.dataset_name,   data_args.dataset_config_name,   split=data_args.train_split_name,   use_auth_token=True if model_args.use_auth_token else None,   )
   UpperCAmelCase__	:Tuple      =      load_dataset(
       data_args.dataset_name,   data_args.dataset_config_name,   split=data_args.eval_split_name,   use_auth_token=True if model_args.use_auth_token else None,   )

   if data_args.audio_column_name not in raw_datasets["train"].column_names:
      raise ValueError(
          F'''--audio_column_name {data_args.audio_column_name} not found in dataset \'{data_args.dataset_name}\'. '''
          '''Make sure to set `--audio_column_name` to the correct audio column - one of '''
          F'''{", ".join(raw_datasets["train"].column_names						)}.'''						)

   if data_args.label_column_name not in raw_datasets["train"].column_names:
      raise ValueError(
          F'''--label_column_name {data_args.label_column_name} not found in dataset \'{data_args.dataset_name}\'. '''
          '''Make sure to set `--label_column_name` to the correct text column - one of '''
          F'''{", ".join(raw_datasets["train"].column_names						)}.'''						)

   # Setting `return_attention_mask=True` is the way to get a correctly masked mean-pooling over
   # transformer outputs in the classifier, but it doesn't always lead to better accuracy
   UpperCAmelCase__	:List[Any]      =      AutoFeatureExtractor.from_pretrained(
       model_args.feature_extractor_name or model_args.model_name_or_path,   return_attention_mask=model_args.attention_mask,   cache_dir=model_args.cache_dir,   revision=model_args.model_revision,   use_auth_token=True if model_args.use_auth_token else None,   )

   # `datasets` takes care of automatically loading and resampling the audio,
   # so we just need to set the correct target sampling rate.
   UpperCAmelCase__	:Optional[Any]      =      raw_datasets.cast_column(
       data_args.audio_column_name,   datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate						)						)

   UpperCAmelCase__	:Any      =      feature_extractor.model_input_names[0]

   def train_transforms(UpperCamelCase_						:      Any						):
      UpperCAmelCase__	:str      =      []
      for audio in batch[data_args.audio_column_name]:
         UpperCAmelCase__	:Any      =      random_subsample(
             audio['''array'''],   max_length=data_args.max_length_seconds,   sample_rate=feature_extractor.sampling_rate						)
         subsampled_wavs.append(UpperCamelCase_						)
      UpperCAmelCase__	:Dict      =      feature_extractor(UpperCamelCase_,   sampling_rate=feature_extractor.sampling_rate						)
      UpperCAmelCase__	:Dict      =      {model_input_name: inputs.get(UpperCamelCase_						)}
      UpperCAmelCase__	:Union[str, Any]      =      list(batch[data_args.label_column_name]						)

      return output_batch

   def val_transforms(UpperCamelCase_						:      Any						):
      UpperCAmelCase__	:str      =      [audio['''array'''] for audio in batch[data_args.audio_column_name]]
      UpperCAmelCase__	:Optional[int]      =      feature_extractor(UpperCamelCase_,   sampling_rate=feature_extractor.sampling_rate						)
      UpperCAmelCase__	:int      =      {model_input_name: inputs.get(UpperCamelCase_						)}
      UpperCAmelCase__	:Any      =      list(batch[data_args.label_column_name]						)

      return output_batch

   # Prepare label mappings.
   # We'll include these in the model's config to get human readable labels in the Inference API.
   UpperCAmelCase__	:List[str]      =      raw_datasets['''train'''].features[data_args.label_column_name].names
   UpperCAmelCase__   ,						UpperCAmelCase__	:List[str]      =      {}, {}
   for i, label in enumerate(UpperCamelCase_						):
      UpperCAmelCase__	:Dict      =      str(UpperCamelCase_						)
      UpperCAmelCase__	:str      =      label

   # Load the accuracy metric from the datasets package
   UpperCAmelCase__	:Any      =      evaluate.load('''accuracy'''						)

   # Define our compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with
   # `predictions` and `label_ids` fields) and has to return a dictionary string to float.
   def compute_metrics(UpperCamelCase_						:      List[str]						):
      UpperCAmelCase__	:List[str]      =      np.argmax(eval_pred.predictions,   axis=1						)
      return metric.compute(predictions=UpperCamelCase_,   references=eval_pred.label_ids						)

   UpperCAmelCase__	:Optional[Any]      =      AutoConfig.from_pretrained(
       model_args.config_name or model_args.model_name_or_path,   num_labels=len(UpperCamelCase_						),   labelaid=UpperCamelCase_,   idalabel=UpperCamelCase_,   finetuning_task='''audio-classification''',   cache_dir=model_args.cache_dir,   revision=model_args.model_revision,   use_auth_token=True if model_args.use_auth_token else None,   )
   UpperCAmelCase__	:List[Any]      =      AutoModelForAudioClassification.from_pretrained(
       model_args.model_name_or_path,   from_tf=bool('''.ckpt''' in model_args.model_name_or_path						),   config=UpperCamelCase_,   cache_dir=model_args.cache_dir,   revision=model_args.model_revision,   use_auth_token=True if model_args.use_auth_token else None,   ignore_mismatched_sizes=model_args.ignore_mismatched_sizes,   )

   # freeze the convolutional waveform encoder
   if model_args.freeze_feature_encoder:
      model.freeze_feature_encoder()

   if training_args.do_train:
      if data_args.max_train_samples is not None:
         UpperCAmelCase__	:List[str]      =      (
             raw_datasets['''train'''].shuffle(seed=training_args.seed						).select(range(data_args.max_train_samples						)						)
         )
      # Set the training transforms
      raw_datasets["train"].set_transform(UpperCamelCase_,   output_all_columns=UpperCamelCase_						)

   if training_args.do_eval:
      if data_args.max_eval_samples is not None:
         UpperCAmelCase__	:List[str]      =      (
             raw_datasets['''eval'''].shuffle(seed=training_args.seed						).select(range(data_args.max_eval_samples						)						)
         )
      # Set the validation transforms
      raw_datasets["eval"].set_transform(UpperCamelCase_,   output_all_columns=UpperCamelCase_						)

   # Initialize our trainer
   UpperCAmelCase__	:List[Any]      =      Trainer(
       model=UpperCamelCase_,   args=UpperCamelCase_,   train_dataset=raw_datasets['''train'''] if training_args.do_train else None,   eval_dataset=raw_datasets['''eval'''] if training_args.do_eval else None,   compute_metrics=UpperCamelCase_,   tokenizer=UpperCamelCase_,   )

   # Training
   if training_args.do_train:
      UpperCAmelCase__	:List[Any]      =      None
      if training_args.resume_from_checkpoint is not None:
         UpperCAmelCase__	:str      =      training_args.resume_from_checkpoint
      elif last_checkpoint is not None:
         UpperCAmelCase__	:int      =      last_checkpoint
      UpperCAmelCase__	:Union[str, Any]      =      trainer.train(resume_from_checkpoint=UpperCamelCase_						)
      trainer.save_model()
      trainer.log_metrics('''train''',   train_result.metrics						)
      trainer.save_metrics('''train''',   train_result.metrics						)
      trainer.save_state()

   # Evaluation
   if training_args.do_eval:
      UpperCAmelCase__	:Optional[int]      =      trainer.evaluate()
      trainer.log_metrics('''eval''',   UpperCamelCase_						)
      trainer.save_metrics('''eval''',   UpperCamelCase_						)

   # Write model card and (optionally) push to hub
   UpperCAmelCase__	:List[Any]      =      {
       '''finetuned_from''': model_args.model_name_or_path,
       '''tasks''': '''audio-classification''',
       '''dataset''': data_args.dataset_name,
       '''tags''': ['''audio-classification'''],
   }
   if training_args.push_to_hub:
      trainer.push_to_hub(**UpperCamelCase_						)
   else:
      trainer.create_model_card(**UpperCamelCase_						)


if __name__ == "__main__":
    main()







'''simple docstring'''





import tempfile
import unittest

from transformers import TaConfig, is_torch_available
from transformers.testing_utils import (
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
    torch_device,
)

from ...generation.test_utils import GenerationTesterMixin
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import AutoTokenizer, UMTaForConditionalGeneration, UMTaForQuestionAnswering, UMTaModel




class 				UpperCAmelCase			:


 def __init__(      self							:	Tuple			,       __lowerCamelCase							:	List[str]			,       __lowerCamelCase							:	Tuple=9_9			,       __lowerCamelCase							:	int=1_3			,       __lowerCamelCase							:	List[str]=7			,       __lowerCamelCase							:	Dict=9			,       __lowerCamelCase							:	List[Any]=True			,       __lowerCamelCase							:	Tuple=True			,       __lowerCamelCase							:	List[str]=False			,       __lowerCamelCase							:	List[Any]=3_2			,       __lowerCamelCase							:	int=5			,       __lowerCamelCase							:	Optional[Any]=4			,       __lowerCamelCase							:	Tuple=3_7			,       __lowerCamelCase							:	Any=8			,       __lowerCamelCase							:	Any=0.1			,       __lowerCamelCase							:	Any=0.0_02			,       __lowerCamelCase							:	List[str]=1			,       __lowerCamelCase							:	str=0			,       __lowerCamelCase							:	str=0			,       __lowerCamelCase							:	str=None			,       __lowerCamelCase							:	List[Any]=None			,       ):
    UpperCAmelCase__	:Tuple      =      parent
    UpperCAmelCase__	:str      =      batch_size
    UpperCAmelCase__	:int      =      encoder_seq_length
    UpperCAmelCase__	:Optional[int]      =      decoder_seq_length
    # For common tests
    UpperCAmelCase__	:int      =      self.decoder_seq_length
    UpperCAmelCase__	:List[Any]      =      is_training
    UpperCAmelCase__	:Any      =      use_attention_mask
    UpperCAmelCase__	:Tuple      =      use_labels
    UpperCAmelCase__	:Optional[int]      =      vocab_size
    UpperCAmelCase__	:Optional[Any]      =      hidden_size
    UpperCAmelCase__	:Optional[Any]      =      num_hidden_layers
    UpperCAmelCase__	:Tuple      =      num_attention_heads
    UpperCAmelCase__	:str      =      d_ff
    UpperCAmelCase__	:Tuple      =      relative_attention_num_buckets
    UpperCAmelCase__	:int      =      dropout_rate
    UpperCAmelCase__	:Dict      =      initializer_factor
    UpperCAmelCase__	:int      =      eos_token_id
    UpperCAmelCase__	:Tuple      =      pad_token_id
    UpperCAmelCase__	:Tuple      =      decoder_start_token_id
    UpperCAmelCase__	:List[str]      =      None
    UpperCAmelCase__	:List[str]      =      decoder_layers


 def        __SCREAMING_SNAKE_CASE					(      self							:	Union[str, Any] ):
    return TaConfig.from_pretrained('''google/umt5-base''' )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[Any]			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	Optional[Any]			,       __lowerCamelCase							:	Optional[int]			,       __lowerCamelCase							:	Optional[Any]=None			,       __lowerCamelCase							:	Dict=None			,       __lowerCamelCase							:	Union[str, Any]=None			,       __lowerCamelCase							:	Optional[int]=None			,       __lowerCamelCase							:	List[Any]=None			,       ):
    if attention_mask is None:
       UpperCAmelCase__	:Optional[Any]      =      input_ids.ne(config.pad_token_id )
    if decoder_attention_mask is None:
       UpperCAmelCase__	:Any      =      decoder_input_ids.ne(config.pad_token_id )
    if head_mask is None:
       UpperCAmelCase__	:List[str]      =      torch.ones(config.num_hidden_layers			,       config.num_attention_heads			,       device=__lowerCamelCase )
    if decoder_head_mask is None:
       UpperCAmelCase__	:Optional[Any]      =      torch.ones(config.num_decoder_layers			,       config.num_attention_heads			,       device=__lowerCamelCase )
    if cross_attn_head_mask is None:
       UpperCAmelCase__	:List[Any]      =      torch.ones(
           config.num_decoder_layers			,       config.num_attention_heads			,       device=__lowerCamelCase )
    return {
        "input_ids": input_ids,
        "decoder_input_ids": decoder_input_ids,
        "attention_mask": attention_mask,
        "decoder_attention_mask": decoder_attention_mask,
        "head_mask": head_mask,
        "decoder_head_mask": decoder_head_mask,
        "cross_attn_head_mask": cross_attn_head_mask,
    }


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    UpperCAmelCase__	:Optional[int]      =      ids_tensor([self.batch_size, self.encoder_seq_length]			,       self.vocab_size )
    UpperCAmelCase__	:List[Any]      =      ids_tensor([self.batch_size, self.decoder_seq_length]			,       self.vocab_size )

    # we need to clamp the input ids here to avoid having pad token in between
    # this is because for NllbMoe the position_ids are prepared such that
    # all pad tokens have pos id = 2 and rest are between 2..seq_length
    # and the seq_length here is seq_length - num_pad_tokens
    # but when using past, there is no way of knowing if the past input ids had
    # pad tokens in them, which results in incorrect seq_lenth and which in turn results in
    # position_ids being off by num_pad_tokens in past input
    UpperCAmelCase__	:int      =      input_ids.clamp(self.pad_token_id + 1 )
    UpperCAmelCase__	:Tuple      =      decoder_input_ids.clamp(self.pad_token_id + 1 )

    UpperCAmelCase__	:Dict      =      self.get_config()
    UpperCAmelCase__	:Union[str, Any]      =      config.num_attention_heads
    UpperCAmelCase__	:Dict      =      self.prepare_inputs_dict(__lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase )
    return config, input_dict


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    UpperCAmelCase__   ,						UpperCAmelCase__	:int      =      self.prepare_config_and_inputs()
    return config, inputs_dict


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    return TaConfig(
        vocab_size=1_6_6			,       d_model=self.hidden_size			,       d_ff=self.d_ff			,       d_kv=self.hidden_size // self.num_attention_heads			,       num_layers=self.num_hidden_layers			,       num_decoder_layers=self.decoder_layers			,       num_heads=self.num_attention_heads			,       relative_attention_num_buckets=self.relative_attention_num_buckets			,       dropout_rate=self.dropout_rate			,       initializer_factor=self.initializer_factor			,       eos_token_id=self.eos_token_id			,       bos_token_id=self.pad_token_id			,       pad_token_id=self.pad_token_id			,       decoder_start_token_id=self.decoder_start_token_id			,       )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[Any] ):
    return TaConfig(
        vocab_size=self.vocab_size			,       d_model=self.hidden_size			,       d_ff=self.d_ff			,       d_kv=self.hidden_size // self.num_attention_heads			,       num_layers=self.num_hidden_layers			,       num_decoder_layers=self.decoder_layers			,       num_heads=self.num_attention_heads			,       relative_attention_num_buckets=self.relative_attention_num_buckets			,       dropout_rate=self.dropout_rate			,       initializer_factor=self.initializer_factor			,       eos_token_id=self.eos_token_id			,       bos_token_id=self.pad_token_id			,       pad_token_id=self.pad_token_id			,       decoder_start_token_id=self.decoder_start_token_id			,       )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Tuple			,       __lowerCamelCase							:	Union[str, Any]			,       __lowerCamelCase							:	List[Any]			,       __lowerCamelCase							:	Optional[int]			,       __lowerCamelCase							:	List[str]			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	List[Any]			,       ):
    UpperCAmelCase__	:str      =      UMTaModel(config=__lowerCamelCase )
    model.to(__lowerCamelCase )
    model.eval()
    UpperCAmelCase__	:Optional[int]      =      model(
        input_ids=__lowerCamelCase			,       decoder_input_ids=__lowerCamelCase			,       attention_mask=__lowerCamelCase			,       decoder_attention_mask=__lowerCamelCase			,       )
    UpperCAmelCase__	:List[Any]      =      model(input_ids=__lowerCamelCase			,       decoder_input_ids=__lowerCamelCase )
    UpperCAmelCase__	:Dict      =      result.last_hidden_state
    UpperCAmelCase__	:Tuple      =      result.past_key_values
    UpperCAmelCase__	:Dict      =      result.encoder_last_hidden_state

    self.parent.assertEqual(encoder_output.size()			,       (self.batch_size, self.encoder_seq_length, self.hidden_size) )
    self.parent.assertEqual(decoder_output.size()			,       (self.batch_size, self.decoder_seq_length, self.hidden_size) )
    # There should be `num_layers` key value embeddings stored in decoder_past
    self.parent.assertEqual(len(__lowerCamelCase )			,       config.num_layers )
    # There should be a self attn key, a self attn value, a cross attn key and a cross attn value stored in each decoder_past tuple
    self.parent.assertEqual(len(decoder_past[0] )			,       4 )


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any]			,       __lowerCamelCase							:	int			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	List[Any]			,       __lowerCamelCase							:	Optional[int]			,       __lowerCamelCase							:	Union[str, Any]			,       __lowerCamelCase							:	List[str]			,       ):
    UpperCAmelCase__	:Tuple      =      UMTaModel(config=__lowerCamelCase ).get_decoder().to(__lowerCamelCase ).eval()
    # first forward pass
    UpperCAmelCase__	:Dict      =      model(__lowerCamelCase			,       use_cache=__lowerCamelCase )
    UpperCAmelCase__	:Any      =      model(__lowerCamelCase )
    UpperCAmelCase__	:List[str]      =      model(__lowerCamelCase			,       use_cache=__lowerCamelCase )

    self.parent.assertTrue(len(__lowerCamelCase ) == len(__lowerCamelCase ) )
    self.parent.assertTrue(len(__lowerCamelCase ) == len(__lowerCamelCase ) + 1 )

    UpperCAmelCase__   ,						UpperCAmelCase__	:int      =      outputs.to_tuple()

    # create hypothetical next token and extent to next_input_ids
    UpperCAmelCase__	:int      =      ids_tensor((self.batch_size, 1)			,       config.vocab_size )

    # append to next input_ids and
    UpperCAmelCase__	:Union[str, Any]      =      torch.cat([input_ids, next_tokens]			,       dim=-1 )

    UpperCAmelCase__	:Union[str, Any]      =      model(__lowerCamelCase )['''last_hidden_state''']
    UpperCAmelCase__	:str      =      model(__lowerCamelCase			,       past_key_values=__lowerCamelCase )['''last_hidden_state''']

    # select random slice
    UpperCAmelCase__	:Tuple      =      ids_tensor((1,)			,       output_from_past.shape[-1] ).item()
    UpperCAmelCase__	:Any      =      output_from_no_past[:, -1, random_slice_idx].detach()
    UpperCAmelCase__	:Union[str, Any]      =      output_from_past[:, 0, random_slice_idx].detach()

    # test that outputs are equal for slice
    self.parent.assertTrue(torch.allclose(__lowerCamelCase			,       __lowerCamelCase			,       atol=1e-3 ) )





 def        __SCREAMING_SNAKE_CASE					(      self							:	Any			,       __lowerCamelCase							:	Tuple			,       __lowerCamelCase							:	str			,       ):
    UpperCAmelCase__	:Optional[Any]      =      UMTaModel(config=__lowerCamelCase ).to(__lowerCamelCase ).half().eval()
    UpperCAmelCase__	:Any      =      model(**__lowerCamelCase )['''last_hidden_state''']
    self.parent.assertFalse(torch.isnan(__lowerCamelCase ).any().item() )




@require_torch
class 				UpperCAmelCase			(			_snake_case						, _snake_case						, _snake_case						, unittest.TestCase						):
 UpperCAmelCase				=					(
     (UMTaModel, UMTaForConditionalGeneration, UMTaForQuestionAnswering) if is_torch_available() else ()
 )
 UpperCAmelCase				=					(UMTaForConditionalGeneration,) if is_torch_available() else ()
 UpperCAmelCase				=					(
     {
         "conversational": UMTaForConditionalGeneration,
         "feature-extraction": UMTaModel,
         "summarization": UMTaForConditionalGeneration,
         "text2text-generation": UMTaForConditionalGeneration,
         "translation": UMTaForConditionalGeneration,
         "question-answering": UMTaForQuestionAnswering,
     }
     if is_torch_available()
     else {}
 )
 UpperCAmelCase				=					True
 UpperCAmelCase				=					False
 UpperCAmelCase				=					False
 UpperCAmelCase				=					True
 UpperCAmelCase				=					True
 # The small UMT5 model needs higher percentages for CPU/MP tests
 UpperCAmelCase				=					[0.8, 0.9]


 def        __SCREAMING_SNAKE_CASE					(      self							:	str ):
    UpperCAmelCase__	:Optional[Any]      =      UMTaModelTester(self )


 @unittest.skip('''Test has a segmentation fault on torch 1.8.0''' )
 def        __SCREAMING_SNAKE_CASE					(      self							:	int ):
    UpperCAmelCase__	:int      =      self.model_tester.prepare_config_and_inputs()
    UpperCAmelCase__	:Optional[int]      =      UMTaModel(config_and_inputs[0] ).to(__lowerCamelCase )
    with tempfile.TemporaryDirectory() as tmpdirname:
       torch.onnx.export(
           __lowerCamelCase			,       (config_and_inputs[1], config_and_inputs[3], config_and_inputs[2])			,       f'''{tmpdirname}/t5_test.onnx'''			,       export_params=__lowerCamelCase			,       opset_version=9			,       input_names=['''input_ids''', '''decoder_input_ids''']			,       )


 @unittest.skipIf(torch_device == '''cpu'''			,       '''Cant do half precision''' )
 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    UpperCAmelCase__	:Optional[int]      =      self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model_fpaa_forward(*__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    UpperCAmelCase__	:Tuple      =      ['''encoder_attentions''', '''decoder_attentions''', '''cross_attentions''']
    UpperCAmelCase__	:Optional[int]      =      self.model_tester.prepare_config_and_inputs()
    UpperCAmelCase__	:int      =      config_and_inputs[0]
    UpperCAmelCase__	:Optional[Any]      =      UMTaForConditionalGeneration(__lowerCamelCase ).eval()
    model.to(__lowerCamelCase )

    UpperCAmelCase__	:List[str]      =      {
        '''head_mask''': torch.zeros(config.num_layers			,       config.num_heads			,       device=__lowerCamelCase ),
        '''decoder_head_mask''': torch.zeros(config.num_decoder_layers			,       config.num_heads			,       device=__lowerCamelCase ),
        '''cross_attn_head_mask''': torch.zeros(config.num_decoder_layers			,       config.num_heads			,       device=__lowerCamelCase ),
    }

    for attn_name, (name, mask) in zip(__lowerCamelCase			,       head_masking.items() ):
       UpperCAmelCase__	:Union[str, Any]      =      {name: mask}
       # Explicitly pass decoder_head_mask as it is required from T5 model when head_mask specified
       if name == "head_mask":
          UpperCAmelCase__	:Optional[Any]      =      torch.ones(
              config.num_decoder_layers			,       config.num_heads			,       device=__lowerCamelCase )

       UpperCAmelCase__	:Union[str, Any]      =      model.generate(
           config_and_inputs[1]['''input_ids''']			,       num_beams=1			,       max_length=3			,       output_attentions=__lowerCamelCase			,       return_dict_in_generate=__lowerCamelCase			,       **__lowerCamelCase			,       )
       # We check the state of decoder_attentions and cross_attentions just from the last step
       UpperCAmelCase__	:List[Any]      =      out[attn_name] if attn_name == attention_names[0] else out[attn_name][-1]
       self.assertEqual(sum([w.sum().item() for w in attn_weights] )			,       0.0 )





 @unittest.skip('''Does not work on the tiny model as we keep hitting edge cases.''' )
 def        __SCREAMING_SNAKE_CASE					(      self							:	Union[str, Any] ):
    pass




@require_torch
@require_sentencepiece
@require_tokenizers
class 				UpperCAmelCase			(			unittest.TestCase						):


 @slow
 @unittest.skip(
     '''Unless we stop stripping left and right by default for all special tokens, the expected ids obtained here will not match the original ones. Wait for https://github.com/huggingface/transformers/pull/23909 to be merged''' )
 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    UpperCAmelCase__	:Optional[Any]      =      UMTaForConditionalGeneration.from_pretrained('''google/umt5-small'''			,       return_dict=__lowerCamelCase ).to(__lowerCamelCase )
    UpperCAmelCase__	:List[str]      =      AutoTokenizer.from_pretrained('''google/umt5-small'''			,       use_fast=__lowerCamelCase			,       legacy=__lowerCamelCase )
    UpperCAmelCase__	:Dict      =      [
        '''Bonjour monsieur <extra_id_0> bien <extra_id_1>.''',
        '''No se como puedo <extra_id_0>.''',
        '''This is the reason why we <extra_id_0> them.''',
        '''The <extra_id_0> walks in <extra_id_1>, seats''',
        '''A <extra_id_0> walks into a bar and orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.''',
    ]
    UpperCAmelCase__	:Optional[Any]      =      tokenizer(__lowerCamelCase			,       return_tensors='''pt'''			,       padding=__lowerCamelCase ).input_ids
    # fmt: off
    UpperCAmelCase__	:List[Any]      =      torch.tensor(
        [
            [ 3_8_5_3_0, 2_1_0_7_0_3, 2_5_6_2_9_9, 1_4_1_0, 2_5_6_2_9_8, 2_7_4, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
            [   8_2_6, 3_2_1, 6_7_1, 2_5_9_2_2, 2_5_6_2_9_9, 2_7_4, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
            [  1_4_6_0, 3_3_9, 3_1_2, 1_9_0_1_4, 1_0_6_2_0, 7_5_8, 2_5_6_2_9_9, 2_3_5_5,2_7_4, 1, 0, 0, 0, 0, 0, 0,0, 0],
            [   5_1_7, 2_5_6_2_9_9, 1_4_8_6_9, 2_8_1, 3_0_1, 2_5_6_2_9_8, 2_7_5, 1_1_9_9_8_3,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
            [   3_2_0, 2_5_6_2_9_9, 1_4_8_6_9, 2_8_1, 2_2_3_4, 2_8_9, 2_2_7_5, 3_3_3,6_1_3_9_1, 2_8_9, 2_5_6_2_9_8, 5_4_3, 2_5_6_2_9_7, 1_6_8_7_1_4, 3_2_9, 2_5_6_2_9_6,2_7_4, 1],
        ] )
    # fmt: on
    torch.testing.assert_allclose(__lowerCamelCase			,       __lowerCamelCase )

    UpperCAmelCase__	:Dict      =      model.generate(input_ids.to(__lowerCamelCase ) )
    UpperCAmelCase__	:str      =      [
        '''<pad><extra_id_0> et<extra_id_1> [eod] <extra_id_2><extra_id_55>.. [eod] 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 <extra_id_56>ajšietosto<extra_id_56>lleux<extra_id_19><extra_id_6>ajšie</s>''',
        '''<pad><extra_id_0>.<extra_id_1>.,<0x0A>...spech <0x0A><extra_id_20> <extra_id_21></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
        '''<pad><extra_id_0> are not going to be a part of the world. We are not going to be a part of<extra_id_1> and<extra_id_2><0x0A><extra_id_48>.<extra_id_48></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
        '''<pad><extra_id_0> door<extra_id_1>, the door<extra_id_2> 피해[/</s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
        '''<pad><extra_id_0>nyone who<extra_id_1> drink<extra_id_2> a<extra_id_3> alcohol<extra_id_4> A<extra_id_5> A. This<extra_id_6> I<extra_id_7><extra_id_52><extra_id_53></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
    ]
    UpperCAmelCase__	:Dict      =      tokenizer.batch_decode(__lowerCamelCase )
    self.assertEqual(__lowerCamelCase			,       __lowerCamelCase )



