

import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, BatchEncoding, PreTrainedTokenizer
from ...utils import logging


UpperCamelCase__  :      str			=    logging.get_logger(__name__)

UpperCamelCase__  :      List[Any]			=    """▁"""

UpperCamelCase__  :      Optional[Any]			=    {"""vocab_file""": """sentencepiece.bpe.model"""}

UpperCamelCase__  :      Any			=    {
    """vocab_file""": {
        """facebook/mbart-large-50-one-to-many-mmt""": (
            """https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/sentencepiece.bpe.model"""
        ),
    }
}

UpperCamelCase__  :      int			=    {
    """facebook/mbart-large-50-one-to-many-mmt""": 1_024,
}

# fmt: off
UpperCamelCase__  :      int			=    ["""ar_AR""", """cs_CZ""", """de_DE""", """en_XX""", """es_XX""", """et_EE""", """fi_FI""", """fr_XX""", """gu_IN""", """hi_IN""", """it_IT""", """ja_XX""", """kk_KZ""", """ko_KR""", """lt_LT""", """lv_LV""", """my_MM""", """ne_NP""", """nl_XX""", """ro_RO""", """ru_RU""", """si_LK""", """tr_TR""", """vi_VN""", """zh_CN""", """af_ZA""", """az_AZ""", """bn_IN""", """fa_IR""", """he_IL""", """hr_HR""", """id_ID""", """ka_GE""", """km_KH""", """mk_MK""", """ml_IN""", """mn_MN""", """mr_IN""", """pl_PL""", """ps_AF""", """pt_XX""", """sv_SE""", """sw_KE""", """ta_IN""", """te_IN""", """th_TH""", """tl_XX""", """uk_UA""", """ur_PK""", """xh_ZA""", """gl_ES""", """sl_SI"""]



class       __snake_case  (    lowerCAmelCase__       ):
  __lowerCAmelCase     :	Optional[int]						 =  VOCAB_FILES_NAMES
  __lowerCAmelCase     :	Union[str, Any]						 =  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  __lowerCAmelCase     :	Optional[int]						 =  PRETRAINED_VOCAB_FILES_MAP
  __lowerCAmelCase     :	Union[str, Any]						 =  ["input_ids", "attention_mask"]

  __lowerCAmelCase     :	List[int]						 =  []
  __lowerCAmelCase     :	List[int]						 =  []

  def __init__(					self						,		_A						,		_A=None						,		_A=None						,		_A="</s>"						,		_A="</s>"						,		_A="<s>"						,		_A="<unk>"						,		_A="<pad>"						,		_A="<mask>"						,		_A = None						,		**_A						,		):
      SCREAMING_SNAKE_CASE_									=						AddedToken(_A						,		lstrip=_A						,		rstrip=_A) if isinstance(_A						,		_A) else mask_token

      SCREAMING_SNAKE_CASE_									=						{} if sp_model_kwargs is None else sp_model_kwargs

      SCREAMING_SNAKE_CASE_									=						kwargs.get('additional_special_tokens'						,		[])
      kwargs["additional_special_tokens"] += [
          code for code in FAIRSEQ_LANGUAGE_CODES if code not in kwargs["additional_special_tokens"]
      ]

      super().__init__(
          src_lang=_A						,		tgt_lang=_A						,		eos_token=_A						,		unk_token=_A						,		sep_token=_A						,		cls_token=_A						,		pad_token=_A						,		mask_token=_A						,		sp_model_kwargs=self.sp_model_kwargs						,		**_A						,		)

      SCREAMING_SNAKE_CASE_									=						spm.SentencePieceProcessor(**self.sp_model_kwargs)
      self.sp_model.Load(str(_A))
      SCREAMING_SNAKE_CASE_									=						vocab_file

      # Original fairseq vocab and spm vocab must be "aligned":
      # Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
      # -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
      # fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
      # spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

      # Mimic fairseq token-to-id alignment for the first 4 token
      SCREAMING_SNAKE_CASE_									=						{'<s>': 0, '<pad>': 1, '</s>': 2, '<unk>': 3}

      # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
      SCREAMING_SNAKE_CASE_									=						1

      SCREAMING_SNAKE_CASE_									=						len(self.sp_model)
      SCREAMING_SNAKE_CASE_									=						{
          code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(_A)
      }
      SCREAMING_SNAKE_CASE_									=						{v: k for k, v in self.lang_code_to_id.items()}
      SCREAMING_SNAKE_CASE_									=						len(self.sp_model) + len(self.lang_code_to_id) + self.fairseq_offset

      self.fairseq_tokens_to_ids.update(self.lang_code_to_id)
      SCREAMING_SNAKE_CASE_									=						{v: k for k, v in self.fairseq_tokens_to_ids.items()}

      SCREAMING_SNAKE_CASE_									=						src_lang if src_lang is not None else 'en_XX'
      SCREAMING_SNAKE_CASE_									=						self.lang_code_to_id[self._src_lang]
      SCREAMING_SNAKE_CASE_									=						tgt_lang
      self.set_src_lang_special_tokens(self._src_lang)

  @property
  def 						lowerCAmelCase__     (					self):
      return len(self.sp_model) + len(self.lang_code_to_id) + self.fairseq_offset + 1  # Plus 1 for the mask token

  @property
  def 						lowerCAmelCase__     (					self):
      return self._src_lang

  @src_lang.setter
  def 						lowerCAmelCase__     (					self						,		_A):
      SCREAMING_SNAKE_CASE_									=						new_src_lang
      self.set_src_lang_special_tokens(self._src_lang)

  def __getstate__(					self):
      SCREAMING_SNAKE_CASE_									=						self.__dict__.copy()
      SCREAMING_SNAKE_CASE_									=						None
      return state

  def __setstate__(					self						,		_A):
      SCREAMING_SNAKE_CASE_									=						d

      # for backward compatibility
      if not hasattr(self						,		'sp_model_kwargs'):
          SCREAMING_SNAKE_CASE_									=						{}

      SCREAMING_SNAKE_CASE_									=						spm.SentencePieceProcessor(**self.sp_model_kwargs)
      self.sp_model.Load(self.vocab_file)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						{self.convert_ids_to_tokens(_A): i for i in range(self.vocab_size)}
      vocab.update(self.added_tokens_encoder)
      return vocab

  def 						lowerCAmelCase__     (					self						,		_A):
      return self.sp_model.encode(_A						,		out_type=_A)

  def 						lowerCAmelCase__     (					self						,		_A):




      if token in self.fairseq_tokens_to_ids:
          return self.fairseq_tokens_to_ids[token]
      SCREAMING_SNAKE_CASE_									=						self.sp_model.PieceToId(_A)

      # Need to return unknown token if the SP model returned 0
      return spm_id + self.fairseq_offset if spm_id else self.unk_token_id

  def 						lowerCAmelCase__     (					self						,		_A):




      if index in self.fairseq_ids_to_tokens:
          return self.fairseq_ids_to_tokens[index]
      return self.sp_model.IdToPiece(index - self.fairseq_offset)

  def 						lowerCAmelCase__     (					self						,		_A):
      SCREAMING_SNAKE_CASE_									=						[]
      SCREAMING_SNAKE_CASE_									=						''
      SCREAMING_SNAKE_CASE_									=						False
      for token in tokens:
          # make sure that special tokens are not decoded using sentencepiece model
          if token in self.all_special_tokens:
              if not prev_is_special:
                  out_string += " "
              out_string += self.sp_model.decode(_A) + token
              SCREAMING_SNAKE_CASE_									=						True
              SCREAMING_SNAKE_CASE_									=						[]
          else:
              current_sub_tokens.append(_A)
              SCREAMING_SNAKE_CASE_									=						False
      out_string += self.sp_model.decode(_A)
      return out_string.strip()

  def 						lowerCAmelCase__     (					self						,		_A						,		_A = None):




      if not os.path.isdir(_A):
          logger.error(f"""Vocabulary path ({save_directory}) should be a directory""")
          return
      SCREAMING_SNAKE_CASE_									=						os.path.join(
          _A						,		(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file'])

      if os.path.abspath(self.vocab_file) != os.path.abspath(_A) and os.path.isfile(self.vocab_file):
          copyfile(self.vocab_file						,		_A)
      elif not os.path.isfile(self.vocab_file):
          with open(_A						,		'wb') as fi:
              SCREAMING_SNAKE_CASE_									=						self.sp_model.serialized_model_proto()
              fi.write(_A)

      return (out_vocab_file,)

  def 						lowerCAmelCase__     (					self						,		_A						,		_A = None						,		_A = False):




      if already_has_special_tokens:
          return super().get_special_tokens_mask(
              token_ids_a=_A						,		token_ids_a=_A						,		already_has_special_tokens=_A)

      SCREAMING_SNAKE_CASE_									=						[1] * len(self.prefix_tokens)
      SCREAMING_SNAKE_CASE_									=						[1] * len(self.suffix_tokens)
      if token_ids_a is None:
          return prefix_ones + ([0] * len(_A)) + suffix_ones
      return prefix_ones + ([0] * len(_A)) + ([0] * len(_A)) + suffix_ones

  def 						lowerCAmelCase__     (					self						,		_A						,		_A = None):




      if token_ids_a is None:
          return self.prefix_tokens + token_ids_a + self.suffix_tokens
      # We don't expect to process pairs, but leave the pair logic for API consistency
      return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens

  def 						lowerCAmelCase__     (					self						,		_A						,		_A						,		_A						,		_A						,		**_A):




      if src_lang is None or tgt_lang is None:
          raise ValueError('Translation requires a `src_lang` and a `tgt_lang` for this model')
      SCREAMING_SNAKE_CASE_									=						src_lang
      SCREAMING_SNAKE_CASE_									=						self(_A						,		add_special_tokens=_A						,		return_tensors=_A						,		**_A)
      SCREAMING_SNAKE_CASE_									=						self.convert_tokens_to_ids(_A)
      SCREAMING_SNAKE_CASE_									=						tgt_lang_id
      return inputs

  def 						lowerCAmelCase__     (					self						,		_A						,		_A = "en_XX"						,		_A = None						,		_A = "ro_RO"						,		**_A						,		):
      SCREAMING_SNAKE_CASE_									=						src_lang
      SCREAMING_SNAKE_CASE_									=						tgt_lang
      return super().prepare_seqaseq_batch(_A						,		_A						,		**_A)

  def 						lowerCAmelCase__     (					self):
      return self.set_src_lang_special_tokens(self.src_lang)

  def 						lowerCAmelCase__     (					self):
      return self.set_tgt_lang_special_tokens(self.tgt_lang)

  def 						lowerCAmelCase__     (					self						,		_A):
      SCREAMING_SNAKE_CASE_									=						self.lang_code_to_id[src_lang]
      SCREAMING_SNAKE_CASE_									=						[self.cur_lang_code_id]
      SCREAMING_SNAKE_CASE_									=						[self.eos_token_id]


  def 						lowerCAmelCase__     (					self						,		_A):
      SCREAMING_SNAKE_CASE_									=						self.lang_code_to_id[tgt_lang]
      SCREAMING_SNAKE_CASE_									=						[self.cur_lang_code_id]
      SCREAMING_SNAKE_CASE_									=						[self.eos_token_id]





import inspect
import os
import unittest
from pathlib import Path

import torch

import accelerate
from accelerate.test_utils import execute_subprocess_async
from accelerate.test_utils.testing import run_command



class       __snake_case  (    unittest.TestCase       ):
  __lowerCAmelCase     :	Dict						 =  inspect.getfile(accelerate.test_utils       )
  __lowerCAmelCase     :	Optional[Any]						 =  os.path.sep.join(mod_file.split(os.path.sep       )[:-1] + ['scripts', 'test_cli.py']       )

  __lowerCAmelCase     :	Tuple						 =  ['accelerate', 'launch']
  __lowerCAmelCase     :	Union[str, Any]						 =  Path.home() / '.cache/huggingface/accelerate'
  __lowerCAmelCase     :	List[str]						 =  'default_config.yaml'
  __lowerCAmelCase     :	List[Any]						 =  config_folder / config_file
  __lowerCAmelCase     :	str						 =  config_folder / '_default_config.yaml'

  __lowerCAmelCase     :	Optional[int]						 =  Path('tests/test_configs'       )

  @classmethod
  def 						lowerCAmelCase__     (					cls):
      if cls.config_path.is_file():
          cls.config_path.rename(cls.changed_path)

  @classmethod
  def 						lowerCAmelCase__     (					cls):
      if cls.changed_path.is_file():
          cls.changed_path.rename(cls.config_path)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.base_cmd
      if torch.cuda.is_available() and (torch.cuda.device_count() > 1):
          cmd += ["--multi_gpu"]
      execute_subprocess_async(cmd + [self.test_file_path]						,		env=os.environ.copy())

  def 						lowerCAmelCase__     (					self):
      for config in sorted(self.test_config_path.glob('**/*.yaml')):
          with self.subTest(config_file=_A):
              execute_subprocess_async(
                  self.base_cmd + ['--config_file', str(_A), self.test_file_path]						,		env=os.environ.copy())


  def 						lowerCAmelCase__     (					self):
      execute_subprocess_async(['accelerate', 'test']						,		env=os.environ.copy())






class       __snake_case  (    unittest.TestCase       ):
  __lowerCAmelCase     :	Optional[Any]						 =  'test-tpu'
  __lowerCAmelCase     :	str						 =  'us-central1-a'
  __lowerCAmelCase     :	Union[str, Any]						 =  'ls'
  __lowerCAmelCase     :	Union[str, Any]						 =  ['accelerate', 'tpu-config']
  __lowerCAmelCase     :	Union[str, Any]						 =  'cd /usr/share'
  __lowerCAmelCase     :	List[Any]						 =  'tests/test_samples/test_command_file.sh'
  __lowerCAmelCase     :	Dict						 =  'Running gcloud compute tpus tpu-vm ssh'

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + ['--command', self.command, '--tpu_zone', self.tpu_zone, '--tpu_name', self.tpu_name, '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/0_12_0.yaml',
              '--command',
              self.command,
              '--tpu_zone',
              self.tpu_zone,
              '--tpu_name',
              self.tpu_name,
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--debug']						,		return_stdout=_A)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--command', self.command, '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/latest.yaml',
              '--command',
              self.command,
              '--command',
              'echo "Hello World"',
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls; echo \"Hello World\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + ['--config_file', 'tests/test_configs/latest.yaml', '--command_file', self.command_file, '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/0_12_0.yaml',
              '--command_file',
              self.command_file,
              '--tpu_zone',
              self.tpu_zone,
              '--tpu_name',
              self.tpu_name,
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--install_accelerate', '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate -U; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/latest.yaml',
              '--install_accelerate',
              '--accelerate_version',
              '12.0.0',
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate==12.0.0; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)


