





"""simple docstring"""






from transformers import BertTokenizer, EncoderDecoderModel, SeqaSeqTrainer, SeqaSeqTrainingArguments
from transformers.testing_utils import TestCasePlus, require_torch, slow
from transformers.utils import is_datasets_available


if is_datasets_available():
    import datasets



class 	lowerCAmelCase					(      __lowerCAmelCase							):






       @slow
       @require_torch
       def 							__A (   self     ):
              _UpperCAmelCase     =     EncoderDecoderModel.from_encoder_decoder_pretrained('prajjwal1/bert-tiny'    ,  'prajjwal1/bert-tiny'     )
              _UpperCAmelCase     =     BertTokenizer.from_pretrained('bert-base-uncased'     )

              _UpperCAmelCase     =     bertabert.config.encoder.vocab_size
              _UpperCAmelCase     =     tokenizer.sep_token_id
              _UpperCAmelCase     =     tokenizer.cls_token_id
              _UpperCAmelCase     =     1_28

              _UpperCAmelCase     =     datasets.load_dataset('cnn_dailymail'    ,  '3.0.0'    ,  split='train[:1%]'     )
              _UpperCAmelCase     =     datasets.load_dataset('cnn_dailymail'    ,  '3.0.0'    ,  split='validation[:1%]'     )

              _UpperCAmelCase     =     train_dataset.select(range(32     )     )
              _UpperCAmelCase     =     val_dataset.select(range(16     )     )

              _UpperCAmelCase     =     4

              def _map_to_encoder_decoder_inputs(a__     ):
                     # Tokenizer will automatically set [BOS] <text> [EOS]
                     _UpperCAmelCase     =     tokenizer(batch['article']    ,  padding='max_length'    ,  truncation=_UpperCamelCase    ,  max_length=5_12     )
                     _UpperCAmelCase     =     tokenizer(batch['highlights']    ,  padding='max_length'    ,  truncation=_UpperCamelCase    ,  max_length=1_28     )
                     _UpperCAmelCase     =     inputs.input_ids
                     _UpperCAmelCase     =     inputs.attention_mask

                     _UpperCAmelCase     =     outputs.input_ids
                     _UpperCAmelCase     =     outputs.input_ids.copy()
                     _UpperCAmelCase     =     [
                         [-1_00 if token == tokenizer.pad_token_id else token for token in labels] for labels in batch["""labels"""]
                     ]
                     _UpperCAmelCase     =     outputs.attention_mask

                     assert all(len(_UpperCamelCase     ) == 5_12 for x in inputs.input_ids     )
                     assert all(len(_UpperCamelCase     ) == 1_28 for x in outputs.input_ids     )

                     return batch

              def _compute_metrics(a__     ):
                     _UpperCAmelCase     =     pred.label_ids
                     _UpperCAmelCase     =     pred.predictions

                     # all unnecessary tokens are removed
                     _UpperCAmelCase     =     tokenizer.batch_decode(_UpperCamelCase    ,  skip_special_tokens=_UpperCamelCase     )
                     _UpperCAmelCase     =     tokenizer.batch_decode(_UpperCamelCase    ,  skip_special_tokens=_UpperCamelCase     )

                     _UpperCAmelCase     =     sum([int(pred_str[i] == label_str[i]     ) for i in range(len(_UpperCamelCase     )     )]     ) / len(_UpperCamelCase     )

                     return {"accuracy": accuracy}

              # map train dataset
              _UpperCAmelCase     =     train_dataset.map(
                  _map_to_encoder_decoder_inputs    ,  batched=_UpperCamelCase    ,  batch_size=_UpperCamelCase    ,  remove_columns=['article', 'highlights']    ,  )
              train_dataset.set_format(
                  type='torch'    ,  columns=['input_ids', 'attention_mask', 'decoder_input_ids', 'decoder_attention_mask', 'labels']    ,  )

              # same for validation dataset
              _UpperCAmelCase     =     val_dataset.map(
                  _map_to_encoder_decoder_inputs    ,  batched=_UpperCamelCase    ,  batch_size=_UpperCamelCase    ,  remove_columns=['article', 'highlights']    ,  )
              val_dataset.set_format(
                  type='torch'    ,  columns=['input_ids', 'attention_mask', 'decoder_input_ids', 'decoder_attention_mask', 'labels']    ,  )

              _UpperCAmelCase     =     self.get_auto_remove_tmp_dir()

              _UpperCAmelCase     =     SeqaSeqTrainingArguments(
                  output_dir=_UpperCamelCase    ,  per_device_train_batch_size=_UpperCamelCase    ,  per_device_eval_batch_size=_UpperCamelCase    ,  predict_with_generate=_UpperCamelCase    ,  evaluation_strategy='steps'    ,  do_train=_UpperCamelCase    ,  do_eval=_UpperCamelCase    ,  warmup_steps=0    ,  eval_steps=2    ,  logging_steps=2    ,  )

              # instantiate trainer
              _UpperCAmelCase     =     SeqaSeqTrainer(
                  model=_UpperCamelCase    ,  args=_UpperCamelCase    ,  compute_metrics=_compute_metrics    ,  train_dataset=_UpperCamelCase    ,  eval_dataset=_UpperCamelCase    ,  tokenizer=_UpperCamelCase    ,  )

              # start training
              trainer.train()













"""simple docstring"""






import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
    import optax
    from flax.training.common_utils import onehot

    from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
    from transformers.models.ta.modeling_flax_ta import shift_tokens_right



@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class 	lowerCAmelCase					(      unittest.TestCase							):






       @slow
       def 							__A (   self     ):
              _UpperCAmelCase     =     FlaxMTaForConditionalGeneration.from_pretrained('google/mt5-small'     )
              _UpperCAmelCase     =     AutoTokenizer.from_pretrained('google/mt5-small'     )

              _UpperCAmelCase     =     tokenizer('Hello there'    ,  return_tensors='np'     ).input_ids
              _UpperCAmelCase     =     tokenizer('Hi I am'    ,  return_tensors='np'     ).input_ids

              _UpperCAmelCase     =     shift_tokens_right(a__    ,  model.config.pad_token_id    ,  model.config.decoder_start_token_id     )

              _UpperCAmelCase     =     model(a__    ,  decoder_input_ids=a__     ).logits
              _UpperCAmelCase     =     optax.softmax_cross_entropy(a__    ,  onehot(a__    ,  logits.shape[-1]     )     ).mean()

              _UpperCAmelCase     =     -(labels.shape[-1] * loss.item())

              _UpperCAmelCase     =     -84.9_127
              self.assertTrue(abs(mtf_score - EXPECTED_SCORE     ) < 1E-4     )






