


'''simple docstring'''





import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


__lowerCamelCase												=						logging.getLogger(__name__)

require_version('''pytorch_lightning>=1.0.4''')

__lowerCamelCase												=						{
    '''base''': AutoModel,
    '''sequence-classification''': AutoModelForSequenceClassification,
    '''question-answering''': AutoModelForQuestionAnswering,
    '''pretraining''': AutoModelForPreTraining,
    '''token-classification''': AutoModelForTokenClassification,
    '''language-modeling''': AutoModelWithLMHead,
    '''summarization''': AutoModelForSeqaSeqLM,
    '''translation''': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
__lowerCamelCase												=						{
    '''linear''': get_linear_schedule_with_warmup,
    '''cosine''': get_cosine_schedule_with_warmup,
    '''cosine_w_restarts''': get_cosine_with_hard_restarts_schedule_with_warmup,
    '''polynomial''': get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
__lowerCamelCase												=						sorted(arg_to_scheduler.keys())
__lowerCamelCase												=						'''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''




class 				UpperCAmelCase			(			pl.LightningModule						):


 def __init__(      self							:	Optional[int]			,       __lowerCamelCase							:	argparse.Namespace			,       __lowerCamelCase							:	Optional[Any]=None			,       __lowerCamelCase							:	Dict="base"			,       __lowerCamelCase							:	Dict=None			,       __lowerCamelCase							:	Dict=None			,       __lowerCamelCase							:	Tuple=None			,       **__lowerCamelCase							:	Tuple			,       ):
    super().__init__()
    # TODO: move to self.save_hyperparameters()
    # self.save_hyperparameters()
    # can also expand arguments into trainer signature for easier reading

    self.save_hyperparameters(__lowerCamelCase )
    UpperCAmelCase__	:List[str]      =      0
    UpperCAmelCase__	:Any      =      Path(self.hparams.output_dir )
    UpperCAmelCase__	:Optional[Any]      =      self.hparams.cache_dir if self.hparams.cache_dir else None
    if config is None:
       UpperCAmelCase__	:Optional[Any]      =      AutoConfig.from_pretrained(
           self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path			,       **({'''num_labels''': num_labels} if num_labels is not None else {})			,       cache_dir=__lowerCamelCase			,       **__lowerCamelCase			,       )
    else:
       UpperCAmelCase__	:PretrainedConfig      =      config

    UpperCAmelCase__	:int      =      ('''encoder_layerdrop''', '''decoder_layerdrop''', '''dropout''', '''attention_dropout''')
    for p in extra_model_params:
       if getattr(self.hparams			,       __lowerCamelCase			,       __lowerCamelCase ):
          assert hasattr(self.config			,       __lowerCamelCase ), f'''model config doesn\'t have a `{p}` attribute'''
          setattr(self.config			,       __lowerCamelCase			,       getattr(self.hparams			,       __lowerCamelCase ) )

    if tokenizer is None:
       UpperCAmelCase__	:List[Any]      =      AutoTokenizer.from_pretrained(
           self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path			,       cache_dir=__lowerCamelCase			,       )
    else:
       UpperCAmelCase__	:PreTrainedTokenizer      =      tokenizer
    UpperCAmelCase__	:List[Any]      =      MODEL_MODES[mode]
    if model is None:
       UpperCAmelCase__	:List[Any]      =      self.model_type.from_pretrained(
           self.hparams.model_name_or_path			,       from_tf=bool('''.ckpt''' in self.hparams.model_name_or_path )			,       config=self.config			,       cache_dir=__lowerCamelCase			,       )
    else:
       UpperCAmelCase__	:List[Any]      =      model


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str]			,       *__lowerCamelCase							:	int			,       **__lowerCamelCase							:	str ):
    UpperCAmelCase__	:Union[str, Any]      =      self.model_type.from_pretrained(*__lowerCamelCase			,       **__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    UpperCAmelCase__	:Optional[int]      =      arg_to_scheduler[self.hparams.lr_scheduler]
    UpperCAmelCase__	:str      =      get_schedule_func(
        self.opt			,       num_warmup_steps=self.hparams.warmup_steps			,       num_training_steps=self.total_steps() )
    UpperCAmelCase__	:Any      =      {'''scheduler''': scheduler, '''interval''': '''step''', '''frequency''': 1}
    return scheduler


 def        __SCREAMING_SNAKE_CASE					(      self							:	str ):
    UpperCAmelCase__	:Union[str, Any]      =      self.model
    UpperCAmelCase__	:int      =      ['''bias''', '''LayerNorm.weight''']
    UpperCAmelCase__	:List[str]      =      [
        {
            '''params''': [
                p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay )
            ],  # check this named paramters
            '''weight_decay''': self.hparams.weight_decay,
        },
        {
            '''params''': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay )],
            '''weight_decay''': 0.0,
        },
    ]
    if self.hparams.adafactor:
       UpperCAmelCase__	:List[Any]      =      Adafactor(
           __lowerCamelCase			,       lr=self.hparams.learning_rate			,       scale_parameter=__lowerCamelCase			,       relative_step=__lowerCamelCase )

    else:
       UpperCAmelCase__	:List[str]      =      AdamW(
           __lowerCamelCase			,       lr=self.hparams.learning_rate			,       eps=self.hparams.adam_epsilon )
    UpperCAmelCase__	:List[Any]      =      optimizer

    UpperCAmelCase__	:Any      =      self.get_lr_scheduler()

    return [optimizer], [scheduler]


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any]			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	List[str] ):
    return self.validation_step(__lowerCamelCase			,       __lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int]			,       __lowerCamelCase							:	Dict ):
    return self.validation_end(__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	int ):
    UpperCAmelCase__	:Any      =      max(1			,       self.hparams.gpus )  # TODO: consider num_tpu_cores
    UpperCAmelCase__	:Dict      =      self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
    return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str]			,       __lowerCamelCase							:	Tuple ):
    if stage == "test":
       UpperCAmelCase__	:str      =      len(self.test_dataloader().dataset )
    else:
       UpperCAmelCase__	:Union[str, Any]      =      self.get_dataloader('''train'''			,       self.hparams.train_batch_size			,       shuffle=__lowerCamelCase )
       UpperCAmelCase__	:Optional[int]      =      len(self.train_dataloader().dataset )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Tuple			,       __lowerCamelCase							:	str			,       __lowerCamelCase							:	int			,       __lowerCamelCase							:	bool = False ):
    raise NotImplementedError('''You must implement this for your task''' )


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any] ):
    return self.train_loader


 def        __SCREAMING_SNAKE_CASE					(      self							:	int ):
    return self.get_dataloader('''dev'''			,       self.hparams.eval_batch_size			,       shuffle=__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    return self.get_dataloader('''test'''			,       self.hparams.eval_batch_size			,       shuffle=__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any]			,       __lowerCamelCase							:	int ):
    return os.path.join(
        self.hparams.data_dir			,       '''cached_{}_{}_{}'''.format(
            __lowerCamelCase			,       list(filter(__lowerCamelCase			,       self.hparams.model_name_or_path.split('''/''' ) ) ).pop()			,       str(self.hparams.max_seq_length )			,       )			,       )


 @pl.utilities.rank_zero_only
 def        __SCREAMING_SNAKE_CASE					(      self							:	Any			,       __lowerCamelCase							:	Dict[str, Any] ):
    UpperCAmelCase__	:Tuple      =      self.output_dir.joinpath('''best_tfmr''' )
    UpperCAmelCase__	:List[Any]      =      self.step_count
    self.model.save_pretrained(__lowerCamelCase )
    self.tokenizer.save_pretrained(__lowerCamelCase )





 @staticmethod
 def        __SCREAMING_SNAKE_CASE					(      __lowerCamelCase							:	Optional[Any]			,       __lowerCamelCase							:	int ):
    parser.add_argument(
        '''--model_name_or_path'''			,       default=__lowerCamelCase			,       type=__lowerCamelCase			,       required=__lowerCamelCase			,       help='''Path to pretrained model or model identifier from huggingface.co/models'''			,       )
    parser.add_argument(
        '''--config_name'''			,       default=''''''			,       type=__lowerCamelCase			,       help='''Pretrained config name or path if not the same as model_name''' )
    parser.add_argument(
        '''--tokenizer_name'''			,       default=__lowerCamelCase			,       type=__lowerCamelCase			,       help='''Pretrained tokenizer name or path if not the same as model_name'''			,       )
    parser.add_argument(
        '''--cache_dir'''			,       default=str(Path(__lowerCamelCase ).parent / '''test_run''' / '''cache''' )			,       type=__lowerCamelCase			,       help='''Where do you want to store the pre-trained models downloaded from huggingface.co'''			,       )
    parser.add_argument(
        '''--encoder_layerdrop'''			,       type=__lowerCamelCase			,       help='''Encoder layer dropout probability (Optional). Goes into model.config'''			,       )
    parser.add_argument(
        '''--decoder_layerdrop'''			,       type=__lowerCamelCase			,       help='''Decoder layer dropout probability (Optional). Goes into model.config'''			,       )
    parser.add_argument(
        '''--dropout'''			,       type=__lowerCamelCase			,       help='''Dropout probability (Optional). Goes into model.config'''			,       )
    parser.add_argument(
        '''--attention_dropout'''			,       type=__lowerCamelCase			,       help='''Attention dropout probability (Optional). Goes into model.config'''			,       )
    parser.add_argument('''--learning_rate'''			,       default=5e-5			,       type=__lowerCamelCase			,       help='''The initial learning rate for Adam.''' )
    parser.add_argument(
        '''--lr_scheduler'''			,       default='''linear'''			,       choices=__lowerCamelCase			,       metavar=__lowerCamelCase			,       type=__lowerCamelCase			,       help='''Learning rate scheduler'''			,       )
    parser.add_argument('''--weight_decay'''			,       default=0.0			,       type=__lowerCamelCase			,       help='''Weight decay if we apply some.''' )
    parser.add_argument('''--adam_epsilon'''			,       default=1e-8			,       type=__lowerCamelCase			,       help='''Epsilon for Adam optimizer.''' )
    parser.add_argument('''--warmup_steps'''			,       default=0			,       type=__lowerCamelCase			,       help='''Linear warmup over warmup_steps.''' )
    parser.add_argument('''--num_workers'''			,       default=4			,       type=__lowerCamelCase			,       help='''kwarg passed to DataLoader''' )
    parser.add_argument('''--num_train_epochs'''			,       dest='''max_epochs'''			,       default=3			,       type=__lowerCamelCase )
    parser.add_argument('''--train_batch_size'''			,       default=3_2			,       type=__lowerCamelCase )
    parser.add_argument('''--eval_batch_size'''			,       default=3_2			,       type=__lowerCamelCase )
    parser.add_argument('''--adafactor'''			,       action='''store_true''' )




class 				UpperCAmelCase			(			pl.Callback						):


 def        __SCREAMING_SNAKE_CASE					(      self							:	str			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	Dict ):
    if (
        trainer.is_global_zero and trainer.global_rank == 0
    ):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
       pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.




class 				UpperCAmelCase			(			pl.Callback						):


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[Any]			,       __lowerCamelCase							:	int			,       __lowerCamelCase							:	List[Any] ):
    # print(pl_module.model.rag)
    for name, param in pl_module.model.rag.named_parameters():
       if param.grad is None:
          print(__lowerCamelCase )




class 				UpperCAmelCase			(			pl.Callback						):


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str]			,       __lowerCamelCase							:	List[str]			,       __lowerCamelCase							:	Any ):
    UpperCAmelCase__	:Optional[int]      =      trainer.lr_schedulers[0]['''scheduler''']
    UpperCAmelCase__	:Optional[int]      =      {f'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr() )}
    pl_module.logger.log_metrics(__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	int			,       __lowerCamelCase							:	pl.Trainer			,       __lowerCamelCase							:	pl.LightningModule ):
    rank_zero_info('''***** Validation results *****''' )
    UpperCAmelCase__	:List[Any]      =      trainer.callback_metrics
    # Log results
    for key in sorted(__lowerCamelCase ):
       if key not in ["log", "progress_bar"]:
          rank_zero_info('''{} = {}\n'''.format(__lowerCamelCase			,       str(metrics[key] ) ) )





 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[Any]			,       __lowerCamelCase							:	pl.Trainer			,       __lowerCamelCase							:	pl.LightningModule ):
    rank_zero_info('''***** Test results *****''' )
    UpperCAmelCase__	:Optional[Any]      =      trainer.callback_metrics
    # Log and save results to file
    UpperCAmelCase__	:Union[str, Any]      =      os.path.join(pl_module.hparams.output_dir			,       '''test_results.txt''' )
    with open(__lowerCamelCase			,       '''w''' ) as writer:
       for key in sorted(__lowerCamelCase ):
          if key not in ["log", "progress_bar"]:
             rank_zero_info('''{} = {}\n'''.format(__lowerCamelCase			,       str(metrics[key] ) ) )
             writer.write('''{} = {}\n'''.format(__lowerCamelCase			,       str(metrics[key] ) ) )


def        a__     (					UpperCamelCase_						:      List[Any],   UpperCamelCase_						:      Tuple						):
   #  To allow all pl args uncomment the following line
   #  parser = pl.Trainer.add_argparse_args(parser)
   parser.add_argument(
       '''--output_dir''',   default=str(Path(UpperCamelCase_						).parent / '''test_run''' / '''model_checkpoints'''						),   type=UpperCamelCase_,   help='''The output directory where the model predictions and checkpoints will be written.''',   )
   parser.add_argument(
       '''--fp16''',   action='''store_true''',   help='''Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit''',   )

   parser.add_argument(
       '''--fp16_opt_level''',   type=UpperCamelCase_,   default='''O2''',   help=(
           '''For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\'].'''
           '''See details at https://nvidia.github.io/apex/amp.html'''
       ),   )
   parser.add_argument('''--n_tpu_cores''',   dest='''tpu_cores''',   type=UpperCamelCase_						)
   parser.add_argument('''--max_grad_norm''',   dest='''gradient_clip_val''',   default=1.0,   type=UpperCamelCase_,   help='''Max gradient norm'''						)
   parser.add_argument('''--do_train''',   action='''store_true''',   help='''Whether to run training.'''						)
   parser.add_argument('''--do_predict''',   action='''store_true''',   help='''Whether to run predictions on the test set.'''						)
   parser.add_argument(
       '''--gradient_accumulation_steps''',   dest='''accumulate_grad_batches''',   type=UpperCamelCase_,   default=1,   help='''Number of updates steps to accumulate before performing a backward/update pass.''',   )
   parser.add_argument('''--seed''',   type=UpperCamelCase_,   default=42,   help='''random seed for initialization'''						)
   parser.add_argument(
       '''--data_dir''',   default=str(Path(UpperCamelCase_						).parent / '''test_run''' / '''dummy-train-data'''						),   type=UpperCamelCase_,   help='''The input data dir. Should contain the training files for the CoNLL-2003 NER task.''',   )


def        a__     (					UpperCamelCase_						:      BaseTransformer,   UpperCamelCase_						:      argparse.Namespace,   UpperCamelCase_						:      int=None,   UpperCamelCase_						:      str=True,   UpperCamelCase_						:      Union[str, Any]=[],   UpperCamelCase_						:      Union[str, Any]=None,   UpperCamelCase_						:      str=None,   **UpperCamelCase_						:      str,   ):
   pl.seed_everything(args.seed						)

   # init model
   UpperCAmelCase__	:Optional[int]      =      Path(model.hparams.output_dir						)
   odir.mkdir(exist_ok=UpperCamelCase_						)

   # add custom checkpoints
   if checkpoint_callback is None:
      UpperCAmelCase__	:List[str]      =      pl.callbacks.ModelCheckpoint(
          filepath=args.output_dir,   prefix='''checkpoint''',   monitor='''val_loss''',   mode='''min''',   save_top_k=1						)
   if early_stopping_callback:
      extra_callbacks.append(UpperCamelCase_						)
   if logging_callback is None:
      UpperCAmelCase__	:Tuple      =      LoggingCallback()

   UpperCAmelCase__	:List[str]      =      {}

   if args.fpaa:
      UpperCAmelCase__	:Optional[Any]      =      16

   if args.gpus > 1:
      UpperCAmelCase__	:Optional[Any]      =      '''auto'''
      UpperCAmelCase__	:str      =      '''ddp'''

   UpperCAmelCase__	:Dict      =      args.accumulate_grad_batches
   UpperCAmelCase__	:Dict      =      None
   UpperCAmelCase__	:Dict      =      '''auto'''

   UpperCAmelCase__	:List[Any]      =      pl.Trainer.from_argparse_args(
       UpperCamelCase_,   weights_summary=UpperCamelCase_,   callbacks=[logging_callback] + extra_callbacks + [InitCallback()] + [checkpoint_callback],   logger=UpperCamelCase_,   val_check_interval=1,   num_sanity_val_steps=2,   **UpperCamelCase_,   )

   if args.do_train:
      trainer.fit(UpperCamelCase_						)

   else:
      print('''RAG modeling tests with new set functions successfuly executed!'''						)
   return trainer







'''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 )



