





import os
from pickle import UnpicklingError
from typing import Dict, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict, unflatten_dict

import transformers

from .utils import logging


__UpperCAmelCase	:       Optional[int]    	=							logging.get_logger(__name__)







def       A__							(   SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__=False)       ->   Tuple:
  try:
    import torch  # noqa: F401
  except ImportError:
    logger.error(
        """Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see"""
        """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
        """ instructions.""")
    raise

  if not is_sharded:
    __snake_case:						Optional[Any]			  =						os.path.abspath(SCREAMING_SNAKE_CASE__)
    logger.info(F'''Loading PyTorch weights from {pt_path}''')

    __snake_case:						List[Any]			  =						torch.load(SCREAMING_SNAKE_CASE__ , map_location="""cpu""")
    logger.info(F'''PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()):,} parameters.''')

    __snake_case:						Optional[int]			  =						convert_pytorch_state_dict_to_flax(SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__)
  else:
    # model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
    __snake_case:						int			  =						convert_pytorch_sharded_state_dict_to_flax(SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__)
  return flax_state_dict







def       A__							(   SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__ , )       ->   (Tuple[str], np.ndarray):

  def is_key_or_prefix_key_in_dict(SCREAMING_SNAKE_CASE__) -> bool:
    return len(set(SCREAMING_SNAKE_CASE__) & {key, (model_prefix,) + key}) > 0

  # layer norm
  __snake_case:						Any			  =						pt_tuple_key[:-1] + ("""scale""",)
  if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(SCREAMING_SNAKE_CASE__):
    return renamed_pt_tuple_key, pt_tensor

  # batch norm layer mean
  __snake_case:						int			  =						pt_tuple_key[:-1] + ("""mean""",)
  if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(SCREAMING_SNAKE_CASE__):
    return renamed_pt_tuple_key, pt_tensor

  # batch norm layer var
  __snake_case:						List[Any]			  =						pt_tuple_key[:-1] + ("""var""",)
  if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(SCREAMING_SNAKE_CASE__):
    return renamed_pt_tuple_key, pt_tensor

  # embedding
  __snake_case:						Any			  =						pt_tuple_key[:-1] + ("""embedding""",)
  if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(SCREAMING_SNAKE_CASE__):
    return renamed_pt_tuple_key, pt_tensor

  # conv layer
  __snake_case:						Tuple			  =						pt_tuple_key[:-1] + ("""kernel""",)
  if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(SCREAMING_SNAKE_CASE__):
    __snake_case:						List[Any]			  =						pt_tensor.transpose(2 , 3 , 1 , 0)
    return renamed_pt_tuple_key, pt_tensor

  # linear layer
  __snake_case:						str			  =						pt_tuple_key[:-1] + ("""kernel""",)
  if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(SCREAMING_SNAKE_CASE__):
    __snake_case:						Optional[int]			  =						pt_tensor.T
    return renamed_pt_tuple_key, pt_tensor

  # old PyTorch layer norm weight
  __snake_case:						str			  =						pt_tuple_key[:-1] + ("""weight""",)
  if pt_tuple_key[-1] == "gamma":
    return renamed_pt_tuple_key, pt_tensor

  # old PyTorch layer norm bias
  __snake_case:						List[Any]			  =						pt_tuple_key[:-1] + ("""bias""",)
  if pt_tuple_key[-1] == "beta":
    return renamed_pt_tuple_key, pt_tensor

  # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
  __snake_case:						Optional[Any]			  =						None
  if pt_tuple_key[-3::2] == ("parametrizations", "original0"):
    __snake_case:						str			  =						pt_tuple_key[-2] + """_g"""
  elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):
    __snake_case:						List[str]			  =						pt_tuple_key[-2] + """_v"""
  if name is not None:
    __snake_case:						Dict			  =						pt_tuple_key[:-3] + (name,)
    return renamed_pt_tuple_key, pt_tensor

  return pt_tuple_key, pt_tensor







def       A__							(   SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__)       ->   Optional[int]:
  # convert pytorch tensor to numpy
  __snake_case:						Any			  =						{k: v.numpy() for k, v in pt_state_dict.items()}

  __snake_case:						List[str]			  =						flax_model.base_model_prefix

  # use params dict if the model contains batch norm layers
  if "params" in flax_model.params:
    __snake_case:						Optional[int]			  =						flax_model.params["""params"""]
  else:
    __snake_case:						List[str]			  =						flax_model.params
  __snake_case:						Union[str, Any]			  =						flatten_dict(SCREAMING_SNAKE_CASE__)

  # add batch_stats keys,values to dict
  if "batch_stats" in flax_model.params:
    __snake_case:						Optional[Any]			  =						flatten_dict(flax_model.params["""batch_stats"""])
    random_flax_state_dict.update(SCREAMING_SNAKE_CASE__)

  __snake_case:						int			  =						{}

  __snake_case:						Dict			  =						(model_prefix not in flax_model_params) and (
      model_prefix in {k.split(""".""")[0] for k in pt_state_dict.keys()}
  )
  __snake_case:						Any			  =						(model_prefix in flax_model_params) and (
      model_prefix not in {k.split(""".""")[0] for k in pt_state_dict.keys()}
  )

  # Need to change some parameters name to match Flax names
  for pt_key, pt_tensor in pt_state_dict.items():
    __snake_case:						int			  =						tuple(pt_key.split("""."""))

    # remove base model prefix if necessary
    __snake_case:						Optional[Any]			  =						pt_tuple_key[0] == model_prefix
    if load_model_with_head_into_base_model and has_base_model_prefix:
      __snake_case:						Optional[int]			  =						pt_tuple_key[1:]

    # Correctly rename weight parameters
    __snake_case  ,       __snake_case:						Any			  =						rename_key_and_reshape_tensor(
        SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__)

    # add model prefix if necessary
    __snake_case:						Optional[int]			  =						(model_prefix,) + flax_key in random_flax_state_dict
    if load_base_model_into_model_with_head and require_base_model_prefix:
      __snake_case:						Any			  =						(model_prefix,) + flax_key

    if flax_key in random_flax_state_dict:
      if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
        raise ValueError(
            F'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
            F'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.''')

        # add batch stats if the model contains batchnorm layers
    if "batch_stats" in flax_model.params:
      if "mean" in flax_key[-1] or "var" in flax_key[-1]:
        __snake_case:						Tuple			  =						jnp.asarray(SCREAMING_SNAKE_CASE__)
        continue
      # remove num_batches_tracked key
      if "num_batches_tracked" in flax_key[-1]:
        flax_state_dict.pop(SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__)
        continue

      # also add unexpected weight so that warning is thrown
      __snake_case:						Dict			  =						jnp.asarray(SCREAMING_SNAKE_CASE__)

    else:
      # also add unexpected weight so that warning is thrown
      __snake_case:						Optional[Any]			  =						jnp.asarray(SCREAMING_SNAKE_CASE__)

  return unflatten_dict(SCREAMING_SNAKE_CASE__)







def       A__							(   SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__)       ->   str:
  import torch

  # Load the index
  __snake_case:						List[str]			  =						{}
  for shard_file in shard_filenames:
    # load using msgpack utils
    __snake_case:						Optional[Any]			  =						torch.load(SCREAMING_SNAKE_CASE__)
    __snake_case:						Union[str, Any]			  =						{k: v.numpy() for k, v in pt_state_dict.items()}

    __snake_case:						Any			  =						flax_model.base_model_prefix

    # use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict
    if "batch_stats" in flax_model.params:
      __snake_case:						Dict			  =						flax_model.params["""params"""]

      __snake_case:						List[str]			  =						flatten_dict(SCREAMING_SNAKE_CASE__)
      random_flax_state_dict.update(flatten_dict(flax_model.params["""batch_stats"""]))
    else:
      __snake_case:						int			  =						flax_model.params
      __snake_case:						Optional[Any]			  =						flatten_dict(SCREAMING_SNAKE_CASE__)

    __snake_case:						Any			  =						(model_prefix not in flax_model_params) and (
        model_prefix in {k.split(""".""")[0] for k in pt_state_dict.keys()}
    )
    __snake_case:						int			  =						(model_prefix in flax_model_params) and (
        model_prefix not in {k.split(""".""")[0] for k in pt_state_dict.keys()}
    )
    # Need to change some parameters name to match Flax names
    for pt_key, pt_tensor in pt_state_dict.items():
      __snake_case:						str			  =						tuple(pt_key.split("""."""))

      # remove base model prefix if necessary
      __snake_case:						List[str]			  =						pt_tuple_key[0] == model_prefix
      if load_model_with_head_into_base_model and has_base_model_prefix:
        __snake_case:						str			  =						pt_tuple_key[1:]

      # Correctly rename weight parameters
      __snake_case  ,       __snake_case:						Union[str, Any]			  =						rename_key_and_reshape_tensor(
          SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__)
      # add model prefix if necessary
      __snake_case:						str			  =						(model_prefix,) + flax_key in random_flax_state_dict
      if load_base_model_into_model_with_head and require_base_model_prefix:
        __snake_case:						Optional[int]			  =						(model_prefix,) + flax_key

      if flax_key in random_flax_state_dict:
        if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
          raise ValueError(
              F'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
              F'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.''')

            # add batch stats if the model contains batchnorm layers
      if "batch_stats" in flax_model.params:
        if "mean" in flax_key[-1]:
          __snake_case:						List[Any]			  =						jnp.asarray(SCREAMING_SNAKE_CASE__)
          continue
        if "var" in flax_key[-1]:
          __snake_case:						Tuple			  =						jnp.asarray(SCREAMING_SNAKE_CASE__)
          continue
        # remove num_batches_tracked key
        if "num_batches_tracked" in flax_key[-1]:
          flax_state_dict.pop(SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__)
          continue

        # also add unexpected weight so that warning is thrown
        __snake_case:						List[Any]			  =						jnp.asarray(SCREAMING_SNAKE_CASE__)

      else:
        # also add unexpected weight so that warning is thrown
        __snake_case:						Tuple			  =						jnp.asarray(SCREAMING_SNAKE_CASE__)
  return unflatten_dict(SCREAMING_SNAKE_CASE__)







def       A__							(   SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__)       ->   int:
  __snake_case:						Dict			  =						os.path.abspath(SCREAMING_SNAKE_CASE__)
  logger.info(F'''Loading Flax weights from {flax_checkpoint_path}''')

  # import correct flax class
  __snake_case:						int			  =						getattr(SCREAMING_SNAKE_CASE__ , """Flax""" + model.__class__.__name__)

  # load flax weight dict
  with open(SCREAMING_SNAKE_CASE__ , """rb""") as state_f:
    try:
      __snake_case:						Union[str, Any]			  =						from_bytes(SCREAMING_SNAKE_CASE__ , state_f.read())
    except UnpicklingError:
      raise EnvironmentError(F'''Unable to convert {flax_checkpoint_path} to Flax deserializable object. ''')

  return load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__)







def       A__							(   SCREAMING_SNAKE_CASE__ , SCREAMING_SNAKE_CASE__)       ->   List[Any]:

  try:
    import torch  # noqa: F401
  except ImportError:
    logger.error(
        """Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see"""
        """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
        """ instructions.""")
    raise

  # check if we have bf16 weights
  __snake_case:						Optional[Any]			  =						flatten_dict(jax.tree_util.tree_map(lambda SCREAMING_SNAKE_CASE__: x.dtype == jnp.bfloataa , SCREAMING_SNAKE_CASE__)).values()
  if any(SCREAMING_SNAKE_CASE__):
    # convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16
    # and bf16 is not fully supported in PT yet.
    logger.warning(
        """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
        """before loading those in PyTorch model.""")
    __snake_case:						List[Any]			  =						jax.tree_util.tree_map(
        lambda SCREAMING_SNAKE_CASE__: params.astype(np.floataa) if params.dtype == jnp.bfloataa else params , SCREAMING_SNAKE_CASE__)

  __snake_case:						Optional[int]			  =						flatten_dict(SCREAMING_SNAKE_CASE__)
  __snake_case:						Tuple			  =						pt_model.state_dict()

  __snake_case:						List[str]			  =						(pt_model.base_model_prefix in flax_state) and (
      pt_model.base_model_prefix not in {k.split(""".""")[0] for k in pt_model_dict.keys()}
  )
  __snake_case:						List[Any]			  =						(pt_model.base_model_prefix not in flax_state) and (
      pt_model.base_model_prefix in {k.split(""".""")[0] for k in pt_model_dict.keys()}
  )

  # keep track of unexpected & missing keys
  __snake_case:						Optional[Any]			  =						[]
  __snake_case:						str			  =						set(pt_model_dict.keys())

  for flax_key_tuple, flax_tensor in flax_state_dict.items():
    __snake_case:						List[str]			  =						flax_key_tuple[0] == pt_model.base_model_prefix
    __snake_case:						int			  =						""".""".join((pt_model.base_model_prefix,) + flax_key_tuple) in pt_model_dict

    # adapt flax_key to prepare for loading from/to base model only
    if load_model_with_head_into_base_model and has_base_model_prefix:
      __snake_case:						str			  =						flax_key_tuple[1:]
    elif load_base_model_into_model_with_head and require_base_model_prefix:
      __snake_case:						str			  =						(pt_model.base_model_prefix,) + flax_key_tuple

    # rename flax weights to PyTorch format
    if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(SCREAMING_SNAKE_CASE__) not in pt_model_dict:
      # conv layer
      __snake_case:						Union[str, Any]			  =						flax_key_tuple[:-1] + ("""weight""",)
      __snake_case:						Dict			  =						jnp.transpose(SCREAMING_SNAKE_CASE__ , (3, 2, 0, 1))
    elif flax_key_tuple[-1] == "kernel" and ".".join(SCREAMING_SNAKE_CASE__) not in pt_model_dict:
      # linear layer
      __snake_case:						str			  =						flax_key_tuple[:-1] + ("""weight""",)
      __snake_case:						Optional[int]			  =						flax_tensor.T
    elif flax_key_tuple[-1] in ["scale", "embedding"]:
      __snake_case:						Union[str, Any]			  =						flax_key_tuple[:-1] + ("""weight""",)

    # adding batch stats from flax batch norm to pt
    elif "mean" in flax_key_tuple[-1]:
      __snake_case:						List[str]			  =						flax_key_tuple[:-1] + ("""running_mean""",)
    elif "var" in flax_key_tuple[-1]:
      __snake_case:						Dict			  =						flax_key_tuple[:-1] + ("""running_var""",)

    if "batch_stats" in flax_state:
      __snake_case:						Any			  =						""".""".join(flax_key_tuple[1:])  # Remove the params/batch_stats header
    else:
      __snake_case:						Dict			  =						""".""".join(SCREAMING_SNAKE_CASE__)

    # We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
    __snake_case:						List[Any]			  =						{}
    # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
    for key in pt_model_dict:
      __snake_case:						List[Any]			  =						key.split(""".""")
      __snake_case:						str			  =						None
      if key_components[-3::2] == ["parametrizations", "original0"]:
        __snake_case:						str			  =						key_components[-2] + """_g"""
      elif key_components[-3::2] == ["parametrizations", "original1"]:
        __snake_case:						Any			  =						key_components[-2] + """_v"""
      if name is not None:
        __snake_case:						int			  =						key_components[:-3] + [name]
        __snake_case:						int			  =						""".""".join(SCREAMING_SNAKE_CASE__)
        __snake_case:						int			  =						key

    if flax_key in special_pt_names:
      __snake_case:						Dict			  =						special_pt_names[flax_key]

    if flax_key in pt_model_dict:
      if flax_tensor.shape != pt_model_dict[flax_key].shape:
        raise ValueError(
            F'''Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '''
            F'''to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.''')
      else:
        # add weight to pytorch dict
        __snake_case:						Dict			  =						np.asarray(SCREAMING_SNAKE_CASE__) if not isinstance(SCREAMING_SNAKE_CASE__ , np.ndarray) else flax_tensor
        __snake_case:						Union[str, Any]			  =						torch.from_numpy(SCREAMING_SNAKE_CASE__)
        # remove from missing keys
        missing_keys.remove(SCREAMING_SNAKE_CASE__)
    else:
      # weight is not expected by PyTorch model
      unexpected_keys.append(SCREAMING_SNAKE_CASE__)

  pt_model.load_state_dict(SCREAMING_SNAKE_CASE__)

  # re-transform missing_keys to list
  __snake_case:						Any			  =						list(SCREAMING_SNAKE_CASE__)

  if len(SCREAMING_SNAKE_CASE__) > 0:
    logger.warning(
        """Some weights of the Flax model were not used when initializing the PyTorch model"""
        F''' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'''
        F''' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'''
        """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
        F''' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'''
        """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
        """ FlaxBertForSequenceClassification model).""")
  else:
    logger.warning(F'''All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n''')
  if len(SCREAMING_SNAKE_CASE__) > 0:
    logger.warning(
        F'''Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'''
        F''' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'''
        """ use it for predictions and inference.""")
  else:
    logger.warning(
        F'''All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n'''
        """If your task is similar to the task the model of the checkpoint was trained on, """
        F'''you can already use {pt_model.__class__.__name__} for predictions without further training.''')

  return pt_model













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


__UpperCAmelCase	:       str    	=							get_tests_dir("fixtures/test_sentencepiece.model")

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

__UpperCAmelCase	:       Any    	=							250_004
__UpperCAmelCase	:       List[str]    	=							250_020







@require_sentencepiece
@require_tokenizers
class    __snake_case	(							__lowerCamelCase       ,      unittest.TestCase      ):





       '''simple docstring'''







       lowerCAmelCase__						=					MBartaaTokenizer
       lowerCAmelCase__						=					MBartaaTokenizerFast
       lowerCAmelCase__						=					True
       lowerCAmelCase__						=					True




       def 							UpperCAmelCase__			(  self			:   Tuple		):
         super().setUp()

         # We have a SentencePiece fixture for testing
         __snake_case:						Optional[int]			  =						MBartaaTokenizer(A							,    src_lang="""en_XX"""							,    tgt_lang="""ro_RO"""							,    keep_accents=A		)
         tokenizer.save_pretrained(self.tmpdirname		)




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         __snake_case:						Any			  =						"""<s>"""
         __snake_case:						Tuple			  =						0

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




       def 							UpperCAmelCase__			(  self			:   Any		):
         __snake_case:						Dict			  =						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(A		)							,    1_054		)




       def 							UpperCAmelCase__			(  self			:   Any		):
         self.assertEqual(self.get_tokenizer().vocab_size							,    1_054		)




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         __snake_case:						Dict			  =						MBartaaTokenizer(A							,    src_lang="""en_XX"""							,    tgt_lang="""ro_RO"""							,    keep_accents=A		)

         __snake_case:						int			  =						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 [285, 46, 10, 170, 382]]							,    )

         __snake_case:						Union[str, Any]			  =						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""", """é""", """."""]							,    )
         __snake_case:						List[Any]			  =						tokenizer.convert_tokens_to_ids(A		)
         self.assertListEqual(
             A							,    [
                 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]
             ]							,    )

         __snake_case:						int			  =						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>""", """."""]							,    )




       @slow
       def 							UpperCAmelCase__			(  self			:   Optional[int]		):
         # fmt: off
         __snake_case:						List[str]			  =						{"""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=A							,    model_name="""facebook/mbart-large-50"""							,    revision="""d3913889c59cd5c9e456b269c376325eabad57e2"""							,    )






       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         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

         __snake_case:						Any			  =						(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})'''		):
             __snake_case:						Optional[int]			  =						self.rust_tokenizer_class.from_pretrained(A							,    **A		)
             __snake_case:						Union[str, Any]			  =						self.tokenizer_class.from_pretrained(A							,    **A		)

             __snake_case:						List[str]			  =						tempfile.mkdtemp()

             __snake_case:						Tuple			  =						tokenizer_r.save_pretrained(A		)
             __snake_case:						Optional[int]			  =						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		)		)
             __snake_case:						Dict			  =						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
             __snake_case:						Tuple			  =						tokenizer_r.from_pretrained(A		)
             __snake_case:						Optional[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
             __snake_case:						Tuple			  =						tempfile.mkdtemp()

             __snake_case:						Any			  =						tokenizer_r.save_pretrained(A							,    legacy_format=A		)
             __snake_case:						List[str]			  =						tokenizer_p.save_pretrained(A		)

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

             # Checks everything loads correctly in the same way
             __snake_case:						List[Any]			  =						tokenizer_r.from_pretrained(A		)
             __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
             __snake_case:						List[str]			  =						tempfile.mkdtemp()

             __snake_case:						Any			  =						tokenizer_r.save_pretrained(A							,    legacy_format=A		)
             __snake_case:						Dict			  =						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
             __snake_case:						Any			  =						tokenizer_r.from_pretrained(A		)
             __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		)		)

             shutil.rmtree(A		)







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





       '''simple docstring'''







       lowerCAmelCase__						=					"""facebook/mbart-large-50-one-to-many-mmt"""
       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__						=					[EN_CODE, 82_74, 12_78_73, 2_59_16, 7, 86_22, 20_71, 4_38, 6_74_85, 53, 18_78_95, 23, 5_17_12, 2]




       @classmethod
       def 							UpperCAmelCase__			(  cls			:   int		):
         __snake_case:						MBartaaTokenizer			  =						MBartaaTokenizer.from_pretrained(
             cls.checkpoint_name							,    src_lang="""en_XX"""							,    tgt_lang="""ro_RO"""		)
         __snake_case:						str			  =						1
         return cls




       def 							UpperCAmelCase__			(  self			:   Any		):
         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 							UpperCAmelCase__			(  self			:   Optional[int]		):
         __snake_case:						List[str]			  =						self.tokenizer.batch_encode_plus(self.src_text		).input_ids[0]
         self.assertListEqual(self.expected_src_tokens							,    A		)




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         self.assertIn(A							,    self.tokenizer.all_special_ids		)
         __snake_case:						Dict			  =						[RO_CODE, 884, 9_019, 96, 9, 916, 86_792, 36, 18_743, 15_596, 5, 2]
         __snake_case:						str			  =						self.tokenizer.decode(A							,    skip_special_tokens=A		)
         __snake_case:						Union[str, Any]			  =						self.tokenizer.decode(generated_ids[1:]							,    skip_special_tokens=A		)
         self.assertEqual(A							,    A		)
         self.assertNotIn(self.tokenizer.eos_token							,    A		)




       def 							UpperCAmelCase__			(  self			:   Dict		):
         __snake_case:						List[str]			  =						["""this is gunna be a long sentence """ * 20]
         assert isinstance(src_text[0]							,    A		)
         __snake_case:						Union[str, Any]			  =						10
         __snake_case:						List[Any]			  =						self.tokenizer(A							,    max_length=A							,    truncation=A		).input_ids[0]
         self.assertEqual(ids[0]							,    A		)
         self.assertEqual(ids[-1]							,    2		)
         self.assertEqual(len(A		)							,    A		)




       def 							UpperCAmelCase__			(  self			:   Tuple		):
         self.assertListEqual(self.tokenizer.convert_tokens_to_ids(["""<mask>""", """ar_AR"""]		)							,    [250_053, 250_001]		)




       def 							UpperCAmelCase__			(  self			:   List[Any]		):
         __snake_case:						List[Any]			  =						tempfile.mkdtemp()
         __snake_case:						Any			  =						self.tokenizer.fairseq_tokens_to_ids
         self.tokenizer.save_pretrained(A		)
         __snake_case:						Union[str, Any]			  =						MBartaaTokenizer.from_pretrained(A		)
         self.assertDictEqual(new_tok.fairseq_tokens_to_ids							,    A		)




       @require_torch
       def 							UpperCAmelCase__			(  self			:   Optional[int]		):
         __snake_case:						List[str]			  =						self.tokenizer(self.src_text							,    text_target=self.tgt_text							,    padding=A							,    return_tensors="""pt"""		)
         __snake_case:						List[Any]			  =						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 							UpperCAmelCase__			(  self			:   Optional[int]		):
         __snake_case:						int			  =						self.tokenizer(
             self.src_text							,    text_target=self.tgt_text							,    padding=A							,    truncation=A							,    max_length=len(self.expected_src_tokens		)							,    return_tensors="""pt"""							,    )
         __snake_case:						Optional[Any]			  =						shift_tokens_right(batch["""labels"""]							,    self.tokenizer.pad_token_id		)

         self.assertIsInstance(A							,    A		)

         self.assertEqual((2, 14)							,    batch.input_ids.shape		)
         self.assertEqual((2, 14)							,    batch.attention_mask.shape		)
         __snake_case:						List[str]			  =						batch.input_ids.tolist()[0]
         self.assertListEqual(self.expected_src_tokens							,    A		)
         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 							UpperCAmelCase__			(  self			:   str		):
         __snake_case:						List[Any]			  =						self.tokenizer(self.src_text							,    padding=A							,    truncation=A							,    max_length=3							,    return_tensors="""pt"""		)
         __snake_case:						Union[str, Any]			  =						self.tokenizer(
             text_target=self.tgt_text							,    padding=A							,    truncation=A							,    max_length=10							,    return_tensors="""pt"""		)
         __snake_case:						Dict			  =						targets["""input_ids"""]
         __snake_case:						Any			  =						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]							,    10		)






       @require_torch
       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         __snake_case:						int			  =						self.tokenizer._build_translation_inputs(
             """A test"""							,    return_tensors="""pt"""							,    src_lang="""en_XX"""							,    tgt_lang="""ar_AR"""		)

         self.assertEqual(
             nested_simplify(A		)							,    {
                 # 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,
             }							,    )






