



import shutil
import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, BatchEncoding, MBartTokenizer, MBartTokenizerFast, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
)

from ...test_tokenization_common import TokenizerTesterMixin


_lowerCamelCase						    =  get_tests_dir("""fixtures/test_sentencepiece.model""")


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

_lowerCamelCase						    =  250004
_lowerCamelCase						    =  250020

@require_sentencepiece
@require_tokenizers
class       _SCREAMING_SNAKE_CASE	(__snake_case			,    unittest.TestCase ):
     lowerCAmelCase					=			MBartTokenizer
     lowerCAmelCase					=			MBartTokenizerFast
     lowerCAmelCase					=			True
     lowerCAmelCase					=			True





     def 	__snake_case    (       self :							Optional[Any]				)->Any:
      super().setUp()

      # We have a SentencePiece fixture for testing
      __SCREAMING_SNAKE_CASE       :    List[str]       =	MBartTokenizer(A_     ,  keep_accents=A_				)
      tokenizer.save_pretrained(self.tmpdirname				)





     def 	__snake_case    (       self :							List[Any]				)->Optional[Any]:
      __SCREAMING_SNAKE_CASE       :    Any       =	MBartTokenizer(A_     ,  keep_accents=A_				)

      __SCREAMING_SNAKE_CASE       :    Tuple       =	tokenizer.tokenize("This is a test"				)
      self.assertListEqual(A_     ,  ["▁This", "▁is", "▁a", "▁t", "est"]				)

      self.assertListEqual(
          tokenizer.convert_tokens_to_ids(A_				)     ,  [value + tokenizer.fairseq_offset for value in [2_8_5, 4_6, 1_0, 1_7_0, 3_8_2]]     ,  )

      __SCREAMING_SNAKE_CASE       :    Tuple       =	tokenizer.tokenize("I was born in 92000, and this is falsé."				)
      self.assertListEqual(
          A_     ,  [
              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",
              "é",
              ".",
          ]     ,  )
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	tokenizer.convert_tokens_to_ids(A_				)
      self.assertListEqual(
          A_     ,  [
              value + tokenizer.fairseq_offset
              for value in [8, 2_1, 8_4, 5_5, 2_4, 1_9, 7, 2, 6_0_2, 3_4_7, 3_4_7, 3_4_7, 3, 1_2, 6_6, 4_6, 7_2, 8_0, 6, 2, 4]
              #                                       ^ unk: 2 + 1 = 3                  unk: 2 + 1 = 3 ^
          ]     ,  )

      __SCREAMING_SNAKE_CASE       :    Tuple       =	tokenizer.convert_ids_to_tokens(A_				)
      self.assertListEqual(
          A_     ,  [
              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>",
              ".",
          ]     ,  )






     def 	__snake_case    (       self :							Optional[int]				)->Optional[int]:




      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

      __SCREAMING_SNAKE_CASE       :    int       =	(self.rust_tokenizer_class, "hf-internal-testing/tiny-random-mbart", {})
      for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
       with self.subTest(F"""{tokenizer.__class__.__name__} ({pretrained_name})"""				):
        __SCREAMING_SNAKE_CASE       :    Dict       =	self.rust_tokenizer_class.from_pretrained(A_     ,  **A_				)
        __SCREAMING_SNAKE_CASE       :    Any       =	self.tokenizer_class.from_pretrained(A_     ,  **A_				)

        __SCREAMING_SNAKE_CASE       :    Any       =	tempfile.mkdtemp()

        __SCREAMING_SNAKE_CASE       :    Any       =	tokenizer_r.save_pretrained(A_				)
        __SCREAMING_SNAKE_CASE       :    Any       =	tokenizer_p.save_pretrained(A_				)

        # 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				)				)
        __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	tuple(f for f in tokenizer_r_files if "tokenizer.json" not in f				)
        self.assertSequenceEqual(A_     ,  A_				)

        # Checks everything loads correctly in the same way
        __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	tokenizer_r.from_pretrained(A_				)
        __SCREAMING_SNAKE_CASE       :    Any       =	tokenizer_p.from_pretrained(A_				)

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

        shutil.rmtree(A_				)

        # Save tokenizer rust, legacy_format=True
        __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	tempfile.mkdtemp()

        __SCREAMING_SNAKE_CASE       :    int       =	tokenizer_r.save_pretrained(A_     ,  legacy_format=A_				)
        __SCREAMING_SNAKE_CASE       :    Tuple       =	tokenizer_p.save_pretrained(A_				)

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

        # Checks everything loads correctly in the same way
        __SCREAMING_SNAKE_CASE       :    int       =	tokenizer_r.from_pretrained(A_				)
        __SCREAMING_SNAKE_CASE       :    Dict       =	tokenizer_p.from_pretrained(A_				)

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

        shutil.rmtree(A_				)

        # Save tokenizer rust, legacy_format=False
        __SCREAMING_SNAKE_CASE       :    List[str]       =	tempfile.mkdtemp()

        __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	tokenizer_r.save_pretrained(A_     ,  legacy_format=A_				)
        __SCREAMING_SNAKE_CASE       :    Any       =	tokenizer_p.save_pretrained(A_				)

        # 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
        __SCREAMING_SNAKE_CASE       :    str       =	tokenizer_r.from_pretrained(A_				)
        __SCREAMING_SNAKE_CASE       :    Optional[int]       =	tokenizer_p.from_pretrained(A_				)

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

        shutil.rmtree(A_				)






@require_torch
@require_sentencepiece
@require_tokenizers
class       _SCREAMING_SNAKE_CASE	(unittest.TestCase ):
     lowerCAmelCase					=			'facebook/mbart-large-en-ro'
     lowerCAmelCase					=			[
         ' 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.',
     ]
     lowerCAmelCase					=			[
         'Ş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.',
     ]
     lowerCAmelCase					=			[8274, 12_7873, 2_5916, 7, 8622, 2071, 438, 6_7485, 53, 18_7895, 23, 5_1712, 2, EN_CODE]





     @classmethod
     def 	__snake_case    (       cls :							Dict				)->Optional[Any]:
      __SCREAMING_SNAKE_CASE       :    MBartTokenizer       =	MBartTokenizer.from_pretrained(
          cls.checkpoint_name     ,  src_lang="en_XX"     ,  tgt_lang="ro_RO"				)
      __SCREAMING_SNAKE_CASE       :    Any       =	1
      return cls





     def 	__snake_case    (       self :							str				)->Union[str, Any]:
      self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["ar_AR"]     ,  2_5_0_0_0_1				)
      self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["en_EN"]     ,  2_5_0_0_0_4				)
      self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["ro_RO"]     ,  2_5_0_0_2_0				)





     def 	__snake_case    (       self :							Dict				)->Dict:
      __SCREAMING_SNAKE_CASE       :    Any       =	self.tokenizer.batch_encode_plus(self.src_text				).input_ids[0]
      self.assertListEqual(self.expected_src_tokens     ,  A_				)





     def 	__snake_case    (       self :							int				)->Any:
      self.assertIn(A_     ,  self.tokenizer.all_special_ids				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	[RO_CODE, 8_8_4, 9_0_1_9, 9_6, 9, 9_1_6, 8_6_7_9_2, 3_6, 1_8_7_4_3, 1_5_5_9_6, 5, 2]
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	self.tokenizer.decode(A_     ,  skip_special_tokens=A_				)
      __SCREAMING_SNAKE_CASE       :    str       =	self.tokenizer.decode(generated_ids[1:]     ,  skip_special_tokens=A_				)
      self.assertEqual(A_     ,  A_				)
      self.assertNotIn(self.tokenizer.eos_token     ,  A_				)





     def 	__snake_case    (       self :							List[Any]				)->List[str]:
      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	["this is gunna be a long sentence " * 2_0]
      assert isinstance(src_text[0]     ,  A_				)
      __SCREAMING_SNAKE_CASE       :    Tuple       =	1_0
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	self.tokenizer(A_     ,  max_length=A_     ,  truncation=A_				).input_ids[0]
      self.assertEqual(ids[-2]     ,  2				)
      self.assertEqual(ids[-1]     ,  A_				)
      self.assertEqual(len(A_				)     ,  A_				)





     def 	__snake_case    (       self :							Optional[int]				)->List[Any]:
      self.assertListEqual(self.tokenizer.convert_tokens_to_ids(["<mask>", "ar_AR"]				)     ,  [2_5_0_0_2_6, 2_5_0_0_0_1]				)





     def 	__snake_case    (       self :							str				)->Optional[Any]:
      __SCREAMING_SNAKE_CASE       :    str       =	tempfile.mkdtemp()
      __SCREAMING_SNAKE_CASE       :    Dict       =	self.tokenizer.fairseq_tokens_to_ids
      self.tokenizer.save_pretrained(A_				)
      __SCREAMING_SNAKE_CASE       :    int       =	MBartTokenizer.from_pretrained(A_				)
      self.assertDictEqual(new_tok.fairseq_tokens_to_ids     ,  A_				)





     @require_torch
     def 	__snake_case    (       self :							Union[str, Any]				)->List[str]:
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	self.tokenizer(self.src_text     ,  text_target=self.tgt_text     ,  padding=A_     ,  return_tensors="pt"				)
      __SCREAMING_SNAKE_CASE       :    Dict       =	shift_tokens_right(batch["labels"]     ,  self.tokenizer.pad_token_id				)

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





     @require_torch
     def 	__snake_case    (       self :							List[Any]				)->int:
      __SCREAMING_SNAKE_CASE       :    Any       =	self.tokenizer(
          self.src_text     ,  text_target=self.tgt_text     ,  padding=A_     ,  truncation=A_     ,  max_length=len(self.expected_src_tokens				)     ,  return_tensors="pt"     ,  )

      __SCREAMING_SNAKE_CASE       :    Dict       =	shift_tokens_right(batch["labels"]     ,  self.tokenizer.pad_token_id				)

      self.assertIsInstance(A_     ,  A_				)

      self.assertEqual((2, 1_4)     ,  batch.input_ids.shape				)
      self.assertEqual((2, 1_4)     ,  batch.attention_mask.shape				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	batch.input_ids.tolist()[0]
      self.assertListEqual(self.expected_src_tokens     ,  A_				)
      self.assertEqual(2     ,  batch.decoder_input_ids[0, -1]				)  # EOS
      # Test that special tokens are reset
      self.assertEqual(self.tokenizer.prefix_tokens     ,  []				)
      self.assertEqual(self.tokenizer.suffix_tokens     ,  [self.tokenizer.eos_token_id, EN_CODE]				)





     def 	__snake_case    (       self :							str				)->Tuple:
      __SCREAMING_SNAKE_CASE       :    int       =	self.tokenizer(self.src_text     ,  padding=A_     ,  truncation=A_     ,  max_length=3     ,  return_tensors="pt"				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	self.tokenizer(
          text_target=self.tgt_text     ,  padding=A_     ,  truncation=A_     ,  max_length=1_0     ,  return_tensors="pt"				)
      __SCREAMING_SNAKE_CASE       :    List[str]       =	targets["input_ids"]
      __SCREAMING_SNAKE_CASE       :    str       =	shift_tokens_right(A_     ,  self.tokenizer.pad_token_id				)

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






     @require_torch
     def 	__snake_case    (       self :							int				)->Tuple:
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	self.tokenizer._build_translation_inputs(
          "A test"     ,  return_tensors="pt"     ,  src_lang="en_XX"     ,  tgt_lang="ar_AR"				)

      self.assertEqual(
          nested_simplify(A_				)     ,  {
              # A, test, EOS, en_XX
              "input_ids": [[6_2, 3_0_3_4, 2, 2_5_0_0_0_4]],
              "attention_mask": [[1, 1, 1, 1]],
              # ar_AR
              "forced_bos_token_id": 2_5_0_0_0_1,
          }     ,  )





from .constants import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    TORCH_LAUNCH_PARAMS,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
)
from .dataclasses import (
    BnbQuantizationConfig,
    ComputeEnvironment,
    CustomDtype,
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    DynamoBackend,
    FPaRecipeKwargs,
    FullyShardedDataParallelPlugin,
    GradientAccumulationPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    KwargsHandler,
    LoggerType,
    MegatronLMPlugin,
    PrecisionType,
    ProjectConfiguration,
    RNGType,
    SageMakerDistributedType,
    TensorInformation,
    TorchDynamoPlugin,
)
from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env
from .imports import (
    get_ccl_version,
    is_abit_bnb_available,
    is_abit_bnb_available,
    is_aim_available,
    is_bfaa_available,
    is_bnb_available,
    is_botoa_available,
    is_ccl_available,
    is_comet_ml_available,
    is_datasets_available,
    is_deepspeed_available,
    is_fpa_available,
    is_ipex_available,
    is_megatron_lm_available,
    is_mlflow_available,
    is_mps_available,
    is_npu_available,
    is_rich_available,
    is_safetensors_available,
    is_sagemaker_available,
    is_tensorboard_available,
    is_tpu_available,
    is_transformers_available,
    is_wandb_available,
    is_xpu_available,
)
from .modeling import (
    check_device_map,
    check_tied_parameters_in_config,
    check_tied_parameters_on_same_device,
    compute_module_sizes,
    convert_file_size_to_int,
    dtype_byte_size,
    find_tied_parameters,
    get_balanced_memory,
    get_max_layer_size,
    get_max_memory,
    get_mixed_precision_context_manager,
    id_tensor_storage,
    infer_auto_device_map,
    load_checkpoint_in_model,
    load_offloaded_weights,
    load_state_dict,
    named_module_tensors,
    retie_parameters,
    set_module_tensor_to_device,
    shard_checkpoint,
)
from .offload import (
    OffloadedWeightsLoader,
    PrefixedDataset,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
    save_offload_index,
)
from .operations import (
    broadcast,
    broadcast_object_list,
    concatenate,
    convert_outputs_to_fpaa,
    convert_to_fpaa,
    find_batch_size,
    find_device,
    gather,
    gather_object,
    get_data_structure,
    honor_type,
    initialize_tensors,
    is_namedtuple,
    is_tensor_information,
    is_torch_tensor,
    listify,
    pad_across_processes,
    recursively_apply,
    reduce,
    send_to_device,
    slice_tensors,
)
from .versions import compare_versions, is_torch_version


if is_deepspeed_available():
      from .deepspeed import (
          DeepSpeedEngineWrapper,
          DeepSpeedOptimizerWrapper,
          DeepSpeedSchedulerWrapper,
          DummyOptim,
          DummyScheduler,
          HfDeepSpeedConfig,
      )

from .bnb import has_abit_bnb_layers, load_and_quantize_model
from .fsdp_utils import load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, save_fsdp_optimizer
from .launch import (
    PrepareForLaunch,
    _filter_args,
    prepare_deepspeed_cmd_env,
    prepare_multi_gpu_env,
    prepare_sagemager_args_inputs,
    prepare_simple_launcher_cmd_env,
    prepare_tpu,
)
from .megatron_lm import (
    AbstractTrainStep,
    BertTrainStep,
    GPTTrainStep,
    MegatronEngine,
    MegatronLMDummyDataLoader,
    MegatronLMDummyScheduler,
    MegatronLMOptimizerWrapper,
    MegatronLMSchedulerWrapper,
    TaTrainStep,
    avg_losses_across_data_parallel_group,
    gather_across_data_parallel_groups,
)
from .megatron_lm import initialize as megatron_lm_initialize
from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
from .megatron_lm import prepare_model as megatron_lm_prepare_model
from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
from .memory import find_executable_batch_size, release_memory
from .other import (
    extract_model_from_parallel,
    get_pretty_name,
    is_port_in_use,
    merge_dicts,
    patch_environment,
    save,
    wait_for_everyone,
    write_basic_config,
)
from .random import set_seed, synchronize_rng_state, synchronize_rng_states
from .torch_xla import install_xla
from .tqdm import tqdm
from .transformer_engine import convert_model, has_transformer_engine_layers






