

import shutil
import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, BatchEncoding, MBartaaTokenizer, MBartaaTokenizerFast, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
)

from ...test_tokenization_common import TokenizerTesterMixin


lowercase_									=   get_tests_dir('fixtures/test_sentencepiece.model')

if is_torch_available():
  from transformers.models.mbart.modeling_mbart import shift_tokens_right

lowercase_									=   2_5_0_0_0_4
lowercase_									=   2_5_0_0_2_0





@require_sentencepiece
@require_tokenizers
class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE			,			unittest.TestCase							):
  _a									=  MBartaaTokenizer
  _a									=  MBartaaTokenizerFast
  _a									=  True
  _a									=  True




  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        super().setUp()

        # We have a SentencePiece fixture for testing
        _lowercase	       =MBartaaTokenizer(lowerCAmelCase							,   src_lang='en_XX'							,   tgt_lang='ro_RO'							,   keep_accents=lowerCAmelCase					)
        tokenizer.save_pretrained(self.tmpdirname					)




  def   A__     (				self					)  -> str:
        '''simple docstring'''






        _lowercase	       ='<s>'
        _lowercase	       =0

        self.assertEqual(self.get_tokenizer()._convert_token_to_id(lowerCAmelCase					)							,   lowerCAmelCase					)
        self.assertEqual(self.get_tokenizer()._convert_id_to_token(lowerCAmelCase					)							,   lowerCAmelCase					)




  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        _lowercase	       =list(self.get_tokenizer().get_vocab().keys()					)

        self.assertEqual(vocab_keys[0]							,   '<s>'					)
        self.assertEqual(vocab_keys[1]							,   '<pad>'					)
        self.assertEqual(vocab_keys[-1]							,   '<mask>'					)
        self.assertEqual(len(lowerCAmelCase					)							,   1_054					)




  def   A__     (				self					)  -> Optional[Any]:
        '''simple docstring'''






        self.assertEqual(self.get_tokenizer().vocab_size							,   1_054					)




  def   A__     (				self					)  -> Tuple:
        '''simple docstring'''






        _lowercase	       =MBartaaTokenizer(lowerCAmelCase							,   src_lang='en_XX'							,   tgt_lang='ro_RO'							,   keep_accents=lowerCAmelCase					)

        _lowercase	       =tokenizer.tokenize('This is a test'					)
        self.assertListEqual(lowerCAmelCase							,   ['▁This', '▁is', '▁a', '▁t', 'est']					)

        self.assertListEqual(
            tokenizer.convert_tokens_to_ids(lowerCAmelCase					)							,   [value + tokenizer.fairseq_offset for value in [285, 46, 10, 170, 382]]							,   )

        _lowercase	       =tokenizer.tokenize('I was born in 92000, and this is falsé.'					)
        self.assertListEqual(
            lowerCAmelCase							,   [SPIECE_UNDERLINE + 'I', SPIECE_UNDERLINE + 'was', SPIECE_UNDERLINE + 'b', 'or', 'n', SPIECE_UNDERLINE + 'in', SPIECE_UNDERLINE + '', '9', '2', '0', '0', '0', ',', SPIECE_UNDERLINE + 'and', SPIECE_UNDERLINE + 'this', SPIECE_UNDERLINE + 'is', SPIECE_UNDERLINE + 'f', 'al', 's', 'é', '.']							,   )
        _lowercase	       =tokenizer.convert_tokens_to_ids(lowerCAmelCase					)
        self.assertListEqual(
            lowerCAmelCase							,   [
                value + tokenizer.fairseq_offset
                for value in [8, 21, 84, 55, 24, 19, 7, 2, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 2, 4]
            ]							,   )

        _lowercase	       =tokenizer.convert_ids_to_tokens(lowerCAmelCase					)
        self.assertListEqual(
            lowerCAmelCase							,   [SPIECE_UNDERLINE + 'I', SPIECE_UNDERLINE + 'was', SPIECE_UNDERLINE + 'b', 'or', 'n', SPIECE_UNDERLINE + 'in', SPIECE_UNDERLINE + '', '<unk>', '2', '0', '0', '0', ',', SPIECE_UNDERLINE + 'and', SPIECE_UNDERLINE + 'this', SPIECE_UNDERLINE + 'is', SPIECE_UNDERLINE + 'f', 'al', 's', '<unk>', '.']							,   )




  @slow
  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        _lowercase	       ={'input_ids': [[250_004, 11_062, 82_772, 7, 15, 82_772, 538, 51_529, 237, 17_198, 1_290, 206, 9, 215_175, 1_314, 136, 17_198, 1_290, 206, 9, 56_359, 42, 122_009, 9, 16_466, 16, 87_344, 4_537, 9, 4_717, 78_381, 6, 159_958, 7, 15, 24_480, 618, 4, 527, 22_693, 5_428, 4, 2_777, 24_480, 9_874, 4, 43_523, 594, 4, 803, 18_392, 33_189, 18, 4, 43_523, 24_447, 12_399, 100, 24_955, 83_658, 9_626, 144_057, 15, 839, 22_335, 16, 136, 24_955, 83_658, 83_479, 15, 39_102, 724, 16, 678, 645, 2_789, 1_328, 4_589, 42, 122_009, 115_774, 23, 805, 1_328, 46_876, 7, 136, 53_894, 1_940, 42_227, 41_159, 17_721, 823, 425, 4, 27_512, 98_722, 206, 136, 5_531, 4_970, 919, 17_336, 5, 2], [250_004, 20_080, 618, 83, 82_775, 47, 479, 9, 1_517, 73, 53_894, 333, 80_581, 110_117, 18_811, 5_256, 1_295, 51, 152_526, 297, 7_986, 390, 124_416, 538, 35_431, 214, 98, 15_044, 25_737, 136, 7_108, 43_701, 23, 756, 135_355, 7, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [250_004, 581, 63_773, 119_455, 6, 147_797, 88_203, 7, 645, 70, 21, 3_285, 10_269, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
        # fmt: on

        self.tokenizer_integration_test_util(
            expected_encoding=lowerCAmelCase							,   model_name='facebook/mbart-large-50'							,   revision='d3913889c59cd5c9e456b269c376325eabad57e2'							,   )




  def   A__     (				self					)  -> str:
        '''simple docstring'''






        if not self.test_slow_tokenizer:
              # as we don't have a slow version, we can't compare the outputs between slow and fast versions
              return

        _lowercase	       =(self.rust_tokenizer_class, 'hf-internal-testing/tiny-random-mbart50', {})
        for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
              with self.subTest(F'''{tokenizer.__class__.__name__} ({pretrained_name})'''					):
                    _lowercase	       =self.rust_tokenizer_class.from_pretrained(lowerCAmelCase							,   **lowerCAmelCase					)
                    _lowercase	       =self.tokenizer_class.from_pretrained(lowerCAmelCase							,   **lowerCAmelCase					)

                    _lowercase	       =tempfile.mkdtemp()

                    _lowercase	       =tokenizer_r.save_pretrained(lowerCAmelCase					)
                    _lowercase	       =tokenizer_p.save_pretrained(lowerCAmelCase					)

                    # Checks it save with the same files + the tokenizer.json file for the fast one
                    self.assertTrue(any('tokenizer.json' in f for f in tokenizer_r_files					)					)
                    _lowercase	       =tuple(f for f in tokenizer_r_files if 'tokenizer.json' not in f					)
                    self.assertSequenceEqual(lowerCAmelCase							,   lowerCAmelCase					)

                    # Checks everything loads correctly in the same way
                    _lowercase	       =tokenizer_r.from_pretrained(lowerCAmelCase					)
                    _lowercase	       =tokenizer_p.from_pretrained(lowerCAmelCase					)

                    # Check special tokens are set accordingly on Rust and Python
                    for key in tokenizer_pp.special_tokens_map:
                          self.assertTrue(hasattr(lowerCAmelCase							,   lowerCAmelCase					)					)
                          # self.assertEqual(getattr(tokenizer_rp, key), getattr(tokenizer_pp, key))
                          # self.assertEqual(getattr(tokenizer_rp, key + "_id"), getattr(tokenizer_pp, key + "_id"))

                    shutil.rmtree(lowerCAmelCase					)

                    # Save tokenizer rust, legacy_format=True
                    _lowercase	       =tempfile.mkdtemp()

                    _lowercase	       =tokenizer_r.save_pretrained(lowerCAmelCase							,   legacy_format=lowerCAmelCase					)
                    _lowercase	       =tokenizer_p.save_pretrained(lowerCAmelCase					)

                    # Checks it save with the same files
                    self.assertSequenceEqual(lowerCAmelCase							,   lowerCAmelCase					)

                    # Checks everything loads correctly in the same way
                    _lowercase	       =tokenizer_r.from_pretrained(lowerCAmelCase					)
                    _lowercase	       =tokenizer_p.from_pretrained(lowerCAmelCase					)

                    # Check special tokens are set accordingly on Rust and Python
                    for key in tokenizer_pp.special_tokens_map:
                          self.assertTrue(hasattr(lowerCAmelCase							,   lowerCAmelCase					)					)

                    shutil.rmtree(lowerCAmelCase					)

                    # Save tokenizer rust, legacy_format=False
                    _lowercase	       =tempfile.mkdtemp()

                    _lowercase	       =tokenizer_r.save_pretrained(lowerCAmelCase							,   legacy_format=lowerCAmelCase					)
                    _lowercase	       =tokenizer_p.save_pretrained(lowerCAmelCase					)

                    # Checks it saved the tokenizer.json file
                    self.assertTrue(any('tokenizer.json' in f for f in tokenizer_r_files					)					)

                    # Checks everything loads correctly in the same way
                    _lowercase	       =tokenizer_r.from_pretrained(lowerCAmelCase					)
                    _lowercase	       =tokenizer_p.from_pretrained(lowerCAmelCase					)

                    # Check special tokens are set accordingly on Rust and Python
                    for key in tokenizer_pp.special_tokens_map:
                          self.assertTrue(hasattr(lowerCAmelCase							,   lowerCAmelCase					)					)

                    shutil.rmtree(lowerCAmelCase					)





@require_torch
@require_sentencepiece
@require_tokenizers
class  __lowerCAmelCase     (					unittest.TestCase							):
  _a									=  """facebook/mbart-large-50-one-to-many-mmt"""
  _a									=  [
      """ UN Chief Says There Is No Military Solution in Syria""",
      """ Secretary-General Ban Ki-moon says his response to Russia's stepped up military support for Syria is that \"there is no military solution\" to the nearly five-year conflict and more weapons will only worsen the violence and misery for millions of people.""",
  ]
  _a									=  [
      """Şeful ONU declară că nu există o soluţie militară în Siria""",
      """Secretarul General Ban Ki-moon declară că răspunsul său la intensificarea sprijinului militar al Rusiei"""
      """ pentru Siria este că \"nu există o soluţie militară\" la conflictul de aproape cinci ani şi că noi arme nu vor"""
      """ face decât să înrăutăţească violenţele şi mizeria pentru milioane de oameni.""",
  ]
  _a									=  [EN_CODE, 8_274, 127_873, 25_916, 7, 8_622, 2_071, 438, 67_485, 53, 187_895, 23, 51_712, 2]




  @classmethod
  def   A__     (				cls					)  -> Optional[Any]:
        '''simple docstring'''






        _lowercase	       =MBartaaTokenizer.from_pretrained(
            cls.checkpoint_name							,   src_lang='en_XX'							,   tgt_lang='ro_RO'					)
        _lowercase	       =1
        return cls




  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['ar_AR']							,   250_001					)
        self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['en_EN']							,   250_004					)
        self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['ro_RO']							,   250_020					)
        self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['mr_IN']							,   250_038					)




  def   A__     (				self					)  -> int:
        '''simple docstring'''






        _lowercase	       =self.tokenizer.batch_encode_plus(self.src_text					).input_ids[0]
        self.assertListEqual(self.expected_src_tokens							,   lowerCAmelCase					)




  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        self.assertIn(lowerCAmelCase							,   self.tokenizer.all_special_ids					)
        _lowercase	       =[RO_CODE, 884, 9_019, 96, 9, 916, 86_792, 36, 18_743, 15_596, 5, 2]
        _lowercase	       =self.tokenizer.decode(lowerCAmelCase							,   skip_special_tokens=lowerCAmelCase					)
        _lowercase	       =self.tokenizer.decode(generated_ids[1:]							,   skip_special_tokens=lowerCAmelCase					)
        self.assertEqual(lowerCAmelCase							,   lowerCAmelCase					)
        self.assertNotIn(self.tokenizer.eos_token							,   lowerCAmelCase					)




  def   A__     (				self					)  -> str:
        '''simple docstring'''






        _lowercase	       =['this is gunna be a long sentence ' * 20]
        assert isinstance(src_text[0]							,   lowerCAmelCase					)
        _lowercase	       =10
        _lowercase	       =self.tokenizer(lowerCAmelCase							,   max_length=lowerCAmelCase							,   truncation=lowerCAmelCase					).input_ids[0]
        self.assertEqual(ids[0]							,   lowerCAmelCase					)
        self.assertEqual(ids[-1]							,   2					)
        self.assertEqual(len(lowerCAmelCase					)							,   lowerCAmelCase					)




  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        self.assertListEqual(self.tokenizer.convert_tokens_to_ids(['<mask>', 'ar_AR']					)							,   [250_053, 250_001]					)




  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        _lowercase	       =tempfile.mkdtemp()
        _lowercase	       =self.tokenizer.fairseq_tokens_to_ids
        self.tokenizer.save_pretrained(lowerCAmelCase					)
        _lowercase	       =MBartaaTokenizer.from_pretrained(lowerCAmelCase					)
        self.assertDictEqual(new_tok.fairseq_tokens_to_ids							,   lowerCAmelCase					)




  @require_torch
  def   A__     (				self					)  -> Any:
        '''simple docstring'''






        _lowercase	       =self.tokenizer(self.src_text							,   text_target=self.tgt_text							,   padding=lowerCAmelCase							,   return_tensors='pt'					)
        _lowercase	       =shift_tokens_right(batch['labels']							,   self.tokenizer.pad_token_id					)

        # fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
        assert batch.input_ids[1][0] == EN_CODE
        assert batch.input_ids[1][-1] == 2
        assert batch.labels[1][0] == RO_CODE
        assert batch.labels[1][-1] == 2
        assert batch.decoder_input_ids[1][:2].tolist() == [2, RO_CODE]




  @require_torch
  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        _lowercase	       =self.tokenizer(
            self.src_text							,   text_target=self.tgt_text							,   padding=lowerCAmelCase							,   truncation=lowerCAmelCase							,   max_length=len(self.expected_src_tokens					)							,   return_tensors='pt'							,   )
        _lowercase	       =shift_tokens_right(batch['labels']							,   self.tokenizer.pad_token_id					)

        self.assertIsInstance(lowerCAmelCase							,   lowerCAmelCase					)

        self.assertEqual((2, 14)							,   batch.input_ids.shape					)
        self.assertEqual((2, 14)							,   batch.attention_mask.shape					)
        _lowercase	       =batch.input_ids.tolist()[0]
        self.assertListEqual(self.expected_src_tokens							,   lowerCAmelCase					)
        self.assertEqual(2							,   batch.decoder_input_ids[0, 0]					)  # decoder_start_token_id
        # Test that special tokens are reset
        self.assertEqual(self.tokenizer.prefix_tokens							,   [EN_CODE]					)
        self.assertEqual(self.tokenizer.suffix_tokens							,   [self.tokenizer.eos_token_id]					)




  def   A__     (				self					)  -> Dict:
        '''simple docstring'''






        _lowercase	       =self.tokenizer(self.src_text							,   padding=lowerCAmelCase							,   truncation=lowerCAmelCase							,   max_length=3							,   return_tensors='pt'					)
        _lowercase	       =self.tokenizer(
            text_target=self.tgt_text							,   padding=lowerCAmelCase							,   truncation=lowerCAmelCase							,   max_length=10							,   return_tensors='pt'					)
        _lowercase	       =targets['input_ids']
        _lowercase	       =shift_tokens_right(lowerCAmelCase							,   self.tokenizer.pad_token_id					)

        self.assertEqual(batch.input_ids.shape[1]							,   3					)
        self.assertEqual(batch.decoder_input_ids.shape[1]							,   10					)




  @require_torch
  def   A__     (				self					)  -> int:
        '''simple docstring'''






        _lowercase	       =self.tokenizer._build_translation_inputs(
            'A test'							,   return_tensors='pt'							,   src_lang='en_XX'							,   tgt_lang='ar_AR'					)

        self.assertEqual(
            nested_simplify(lowerCAmelCase					)							,   {
                # en_XX, A, test, EOS
                'input_ids': [[250_004, 62, 3_034, 2]],
                'attention_mask': [[1, 1, 1, 1]],
                # ar_AR
                'forced_bos_token_id': 250_001,
            }							,   )





import math
import os
import re
import sys
import unittest
from pathlib import Path
from typing import Tuple
from unittest.mock import patch

from parameterized import parameterized

from transformers.testing_utils import (
    CaptureStderr,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    get_torch_dist_unique_port,
    require_apex,
    require_bitsandbytes,
    require_fairscale,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    require_torch_non_multi_gpu,
    slow,
)
from transformers.trainer_callback import TrainerState
from transformers.trainer_utils import set_seed


lowercase_									=   os.path.abspath(os.path.dirname(__file__))
with ExtendSysPath(f"{bindir}/../../examples/pytorch/translation"):
  from run_translation import main  # noqa


set_seed(4_2)
lowercase_									=   'sshleifer/student_marian_en_ro_6_1'
lowercase_									=   'sshleifer/tiny-mbart'





@require_torch
class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE							):




  def   A__     (				self							,   lowerCAmelCase=False							,   lowerCAmelCase=None							,   lowerCAmelCase=True							,   lowerCAmelCase=True							,   lowerCAmelCase=True							,   lowerCAmelCase=True							,   )  -> Dict:
        '''simple docstring'''






        _lowercase	       =self.run_trainer(
            eval_steps=1							,   max_len=12							,   model_name=lowerCAmelCase							,   num_train_epochs=1							,   distributed=lowerCAmelCase							,   extra_args_str=lowerCAmelCase							,   predict_with_generate=lowerCAmelCase							,   do_train=lowerCAmelCase							,   do_eval=lowerCAmelCase							,   do_predict=lowerCAmelCase							,   )
        _lowercase	       =TrainerState.load_from_json(os.path.join(lowerCAmelCase							,   'trainer_state.json'					)					).log_history

        if not do_eval:
              return

        _lowercase	       =[log for log in logs if 'eval_loss' in log.keys()]

        _lowercase	       =eval_metrics[0]
        if predict_with_generate:
              assert "eval_bleu" in first_step_stats

              _lowercase	       =eval_metrics[-1]
              assert isinstance(last_step_stats['eval_bleu']							,   lowerCAmelCase					)
              assert not math.isnan(float(last_step_stats['eval_loss']					)					), "eval_loss must not be `nan`"




  @require_torch_non_multi_gpu
  def   A__     (				self					)  -> Optional[Any]:
        '''simple docstring'''






        self.run_seqaseq_quick()




  @require_torch_multi_gpu
  def   A__     (				self					)  -> Tuple:
        '''simple docstring'''






        self.run_seqaseq_quick(distributed=lowerCAmelCase					)




  @require_torch_multi_gpu
  def   A__     (				self					)  -> Optional[Any]:
        '''simple docstring'''






        self.run_seqaseq_quick(distributed=lowerCAmelCase					)




  @unittest.skip('Requires an update of the env running those tests'					)
  @require_torch_multi_gpu
  @require_fairscale
  def   A__     (				self					)  -> Optional[Any]:
        '''simple docstring'''






        self.run_seqaseq_quick(distributed=lowerCAmelCase							,   extra_args_str='--sharded_ddp simple'					)




  @unittest.skip('Requires an update of the env running those tests'					)
  @require_torch_multi_gpu
  @require_fairscale
  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        self.run_seqaseq_quick(distributed=lowerCAmelCase							,   extra_args_str='--sharded_ddp simple --fp16'					)




  @unittest.skip('Requires an update of the env running those tests'					)
  @require_torch_multi_gpu
  @require_fairscale
  def   A__     (				self					)  -> Dict:
        '''simple docstring'''






        self.run_seqaseq_quick(distributed=lowerCAmelCase							,   extra_args_str='--sharded_ddp zero_dp_2'							,   predict_with_generate=lowerCAmelCase					)




  @unittest.skip('Requires an update of the env running those tests'					)
  @require_torch_multi_gpu
  @require_fairscale
  def   A__     (				self					)  -> Optional[int]:
        '''simple docstring'''






        self.run_seqaseq_quick(
            distributed=lowerCAmelCase							,   extra_args_str='--sharded_ddp zero_dp_2 --fp16'							,   predict_with_generate=lowerCAmelCase					)




  @require_apex
  @require_torch_gpu
  def   A__     (				self					)  -> List[Any]:
        '''simple docstring'''






        self.run_seqaseq_quick(distributed=lowerCAmelCase							,   extra_args_str='--fp16 --fp16_backend=apex'					)
        # test 2nd time - was getting eval_loss': nan'
        # to reproduce the problem set distributed=False
        self.run_seqaseq_quick(distributed=lowerCAmelCase							,   extra_args_str='--fp16 --fp16_backend=apex'					)




  @parameterized.expand(['base', 'low', 'high', 'mixed']					)
  @require_torch_multi_gpu
  def   A__     (				self							,   lowerCAmelCase					)  -> Any:
        '''simple docstring'''






        _lowercase	       ={
            # test with the default log_level - should be info and thus log info once
            'base': {'extra_args_str': '', 'n_matches': 1},
            # test with low log_level and log_level_replica - should be noisy on all processes
            # now the info string should appear twice on 2 processes
            'low': {'extra_args_str': '--log_level debug --log_level_replica debug', 'n_matches': 2},
            # test with high log_level and low log_level_replica
            # now the info string should appear once only on the replica
            'high': {'extra_args_str': '--log_level error --log_level_replica debug', 'n_matches': 1},
            # test with high log_level and log_level_replica - should be quiet on all processes
            'mixed': {'extra_args_str': '--log_level error --log_level_replica error', 'n_matches': 0},
        }

        _lowercase	       =experiments[experiment_id]
        _lowercase	       ={'distributed': True, 'predict_with_generate': False, 'do_eval': False, 'do_predict': False}
        _lowercase	       ='Running training'
        with CaptureStderr() as cl:
              self.run_seqaseq_quick(**lowerCAmelCase							,   extra_args_str=data['extra_args_str']					)
        _lowercase	       =len(re.findall(lowerCAmelCase							,   cl.err					)					)
        self.assertEqual(lowerCAmelCase							,   data['n_matches']					)




  @slow
  def   A__     (				self					)  -> Union[str, Any]:
        '''simple docstring'''






        _lowercase	       =self.run_trainer(
            eval_steps=2							,   max_len=128							,   model_name=lowerCAmelCase							,   learning_rate=3e-4							,   num_train_epochs=10							,   distributed=lowerCAmelCase							,   )

        # Check metrics
        _lowercase	       =TrainerState.load_from_json(os.path.join(lowerCAmelCase							,   'trainer_state.json'					)					).log_history
        _lowercase	       =[log for log in logs if 'eval_loss' in log.keys()]
        _lowercase	       =eval_metrics[0]
        _lowercase	       =eval_metrics[-1]

        assert first_step_stats["eval_loss"] > last_step_stats["eval_loss"], "model learned nothing"
        assert isinstance(last_step_stats['eval_bleu']							,   lowerCAmelCase					)

        # test if do_predict saves generations and metrics
        _lowercase	       =os.listdir(lowerCAmelCase					)
        _lowercase	       ={os.path.basename(lowerCAmelCase					) for p in contents}
        assert "generated_predictions.txt" in contents
        assert "predict_results.json" in contents




  @slow
  @require_bitsandbytes
  def   A__     (				self					)  -> List[str]:
        '''simple docstring'''






        from transformers.training_args import OptimizerNames

        def train_and_return_metrics(lowerCAmelCase					) -> Tuple[int, float]:
              _lowercase	       ='--skip_memory_metrics 0'

              _lowercase	       =self.run_trainer(
                  max_len=128							,   model_name=lowerCAmelCase							,   learning_rate=3e-4							,   num_train_epochs=1							,   optim=lowerCAmelCase							,   distributed=lowerCAmelCase							,   extra_args_str=lowerCAmelCase							,   do_eval=lowerCAmelCase							,   do_predict=lowerCAmelCase							,   n_gpus_to_use=1							,   )

              # Check metrics
              _lowercase	       =TrainerState.load_from_json(Path(lowerCAmelCase							,   'trainer_state.json'					)					).log_history
              _lowercase	       =int(logs[0]['train_mem_gpu_peaked_delta'] / 2**20					)
              _lowercase	       =int(logs[0]['train_mem_gpu_alloc_delta'] / 2**20					)

              _lowercase	       =logs[0]['train_loss']
              return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

        _lowercase       ,       _lowercase       ,       _lowercase	       =train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value					)
        _lowercase       ,       _lowercase       ,       _lowercase	       =train_and_return_metrics(OptimizerNames.ADAMW_BNB.value					)

        _lowercase	       =gpu_alloc_mem_orig - gpu_alloc_mem_bnb

        _lowercase	       =gpu_peak_mem_orig + gpu_alloc_mem_orig
        _lowercase	       =gpu_peak_mem_bnb + gpu_alloc_mem_bnb
        _lowercase	       =gpu_total_mem_orig - gpu_total_mem_bnb

        # sshleifer/student_marian_en_ro_6_1 has 54M parameter, 29M of which is `nn.Embedding` which
        # doesn't get quantized and remains in fp32. Therefore we only have 25M parameters quantized
        # in 2 bytes and the diff in optim memory usage is derived as so:
        #
        # - normal 25*8=~200MB (8 bytes per param)
        # - bnb    25*2= ~50MB (2 bytes per param)
        #
        # Thus we should expect ~150MB total memory saved.
        #
        # Peak memory should be the same - the total should be different by about that same margin
        #
        # After leaving a small margin to accommodate for differences between gpus let's check
        # that we have at least 120MB in savings
        _lowercase	       =120

        # uncomment the following if this test starts failing - requires py38 for a new print feature
        # gpu_peak_mem_diff = gpu_peak_mem_orig - gpu_peak_mem_bnb
        # print(f"{gpu_alloc_mem_orig=}MB {gpu_peak_mem_orig=}MB {gpu_alloc_mem_orig+gpu_peak_mem_orig=}MB")
        # print(f" {gpu_alloc_mem_bnb=}MB  {gpu_peak_mem_bnb=}MB  {gpu_alloc_mem_bnb+gpu_peak_mem_bnb=}MB")
        # print(f"{gpu_alloc_mem_diff=}MB")
        # print(f"{gpu_peak_mem_diff=}MB")
        # print(f"{gpu_total_mem_orig=}MB, {gpu_total_mem_bnb=}MB")
        # print(f"{gpu_total_mem_diff=}MB, {gpu_total_mem_diff=}MB")

        self.assertGreater(
            lowerCAmelCase							,   lowerCAmelCase							,   'should use ~150MB less alloc gpu memory with BNB, compared to without it for this model but got'
            F''' a difference of {gpu_alloc_mem_diff}MB, with gpu_alloc_mem_orig={gpu_alloc_mem_orig}MB and'''
            F''' gpu_alloc_mem_bnb={gpu_alloc_mem_bnb}MB'''							,   )

        self.assertGreater(
            lowerCAmelCase							,   lowerCAmelCase							,   'should use ~150MB less total gpu memory with BNB, compared to without it for this model but got'
            F''' a difference of {gpu_total_mem_diff}MB, with gpu_total_mem_orig={gpu_total_mem_orig}MB and'''
            F''' gpu_total_mem_bnb={gpu_total_mem_bnb}MB'''							,   )

        self.assertEqual(
            lowerCAmelCase							,   lowerCAmelCase							,   F'''loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}'''					)




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase = 3e-3							,   lowerCAmelCase = "adafactor"							,   lowerCAmelCase = False							,   lowerCAmelCase = None							,   lowerCAmelCase = 0							,   lowerCAmelCase = True							,   lowerCAmelCase = True							,   lowerCAmelCase = True							,   lowerCAmelCase = True							,   lowerCAmelCase = None							,   )  -> Optional[Any]:
        '''simple docstring'''






        _lowercase	       =self.test_file_dir / '../fixtures/tests_samples/wmt_en_ro'
        _lowercase	       =self.get_auto_remove_tmp_dir()
        _lowercase	       =F'''
            --model_name_or_path {model_name}
            --train_file {data_dir}/train.json
            --validation_file {data_dir}/val.json
            --test_file {data_dir}/test.json
            --output_dir {output_dir}
            --overwrite_output_dir
            --max_train_samples 8
            --max_source_length {max_len}
            --max_target_length {max_len}
            --do_train
            --num_train_epochs {str(lowerCAmelCase					)}
            --per_device_train_batch_size 4
            --learning_rate {learning_rate}
            --warmup_steps 8
            --logging_steps 0
            --logging_strategy no
            --save_steps {str(lowerCAmelCase					)}
            --group_by_length
            --label_smoothing_factor 0.1
            --target_lang ro_RO
            --source_lang en_XX
        '''.split()

        _lowercase	       =F'''
            --do_eval
            --per_device_eval_batch_size 4
            --max_eval_samples 8
            --val_max_target_length {max_len}
            --evaluation_strategy steps
            --eval_steps {str(lowerCAmelCase					)}
        '''.split()

        _lowercase	       ='\n            --do_predict\n        '.split()

        _lowercase	       =[]
        if do_train:
              args += args_train

        if do_eval:
              args += args_eval

        if do_predict:
              args += args_predict

        if predict_with_generate:
              args += "--predict_with_generate".split()

        if do_train:
              if optim == "adafactor":
                    args += "--adafactor".split()
              else:
                    args += F'''--optim {optim}'''.split()

        if extra_args_str is not None:
              args += extra_args_str.split()

        if distributed:
              if n_gpus_to_use is None:
                    _lowercase	       =get_gpu_count()
              _lowercase	       =get_torch_dist_unique_port()
              _lowercase	       =F'''
                -m torch.distributed.run
                --nproc_per_node={n_gpus_to_use}
                --master_port={master_port}
                {self.examples_dir_str}/pytorch/translation/run_translation.py
            '''.split()
              _lowercase	       =[sys.executable] + distributed_args + args
              # keep for quick debug
              # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
              execute_subprocess_async(lowerCAmelCase							,   env=self.get_env()					)
        else:
              _lowercase	       =['run_translation.py'] + args
              with patch.object(lowerCAmelCase							,   'argv'							,   lowerCAmelCase					):
                    main()

        return output_dir


