






"""simple docstring"""






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

from tokenizers import processors

from ...tokenization_utils import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
  from .tokenization_mbart import MBartTokenizer
else:
  UpperCAmelCase_   : Any   			= None


UpperCAmelCase_   : Dict   			= logging.get_logger(__name__)


UpperCAmelCase_   : List[Any]   			= {"""vocab_file""": """sentencepiece.bpe.model""", """tokenizer_file""": """tokenizer.json"""}

UpperCAmelCase_   : List[Any]   			= {
    """vocab_file""": {
        """facebook/mbart-large-en-ro""": (
            """https://huggingface.co/facebook/mbart-large-en-ro/resolve/main/sentencepiece.bpe.model"""
        ),
        """facebook/mbart-large-cc25""": (
            """https://huggingface.co/facebook/mbart-large-cc25/resolve/main/sentencepiece.bpe.model"""
        ),
    },
    """tokenizer_file""": {
        """facebook/mbart-large-en-ro""": """https://huggingface.co/facebook/mbart-large-en-ro/resolve/main/tokenizer.json""",
        """facebook/mbart-large-cc25""": """https://huggingface.co/facebook/mbart-large-cc25/resolve/main/tokenizer.json""",
    },
}

UpperCAmelCase_   : List[Any]   			= {
    """facebook/mbart-large-en-ro""": 1024,
    """facebook/mbart-large-cc25""": 1024,
}

# fmt: off
UpperCAmelCase_   : List[Any]   			= ["""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"""]




class 							lowerCAmelCase__				(			UpperCAmelCase__ ):






      '''simple docstring'''

      __UpperCamelCase	 =  VOCAB_FILES_NAMES
      __UpperCamelCase	 =  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __UpperCamelCase	 =  PRETRAINED_VOCAB_FILES_MAP
      __UpperCamelCase	 =  ["input_ids", "attention_mask"]
      __UpperCamelCase	 =  MBartTokenizer

      __UpperCamelCase	 =  []
      __UpperCamelCase	 =  []
      def __init__(						self     :  Any						,	lowercase_     :  str=None						,	lowercase_     :  int=None						,	lowercase_     :  List[Any]="<s>"						,	lowercase_     :  Any="</s>"						,	lowercase_     :  Any="</s>"						,	lowercase_     :  List[str]="<s>"						,	lowercase_     :  List[Any]="<unk>"						,	lowercase_     :  List[Any]="<pad>"						,	lowercase_     :  List[str]="<mask>"						,	lowercase_     :  Optional[int]=None						,	lowercase_     :  int=None						,	lowercase_     :  List[Any]=None						,	**lowercase_     :  str						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Dict    						=		AddedToken(lowercase_						,	lstrip=lowercase_						,	rstrip=lowercase_) if isinstance(lowercase_						,	lowercase_) else mask_token

          super().__init__(
              vocab_file=lowercase_						,	tokenizer_file=lowercase_						,	bos_token=lowercase_						,	eos_token=lowercase_						,	sep_token=lowercase_						,	cls_token=lowercase_						,	unk_token=lowercase_						,	pad_token=lowercase_						,	mask_token=lowercase_						,	src_lang=lowercase_						,	tgt_lang=lowercase_						,	additional_special_tokens=lowercase_						,	**lowercase_						,	)

          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		vocab_file
          SCREAMING_SNAKE_CASE_      :  Any    						=		False if not self.vocab_file else True

          SCREAMING_SNAKE_CASE_      :  Dict    						=		FAIRSEQ_LANGUAGE_CODES.copy()

          if additional_special_tokens is not None:
              # Only add those special tokens if they are not already there.
              _additional_special_tokens.extend(
                  [t for t in additional_special_tokens if t not in _additional_special_tokens])

          self.add_special_tokens({'''additional_special_tokens''': _additional_special_tokens})
          SCREAMING_SNAKE_CASE_      :  Any    						=		{
              lang_code: self.convert_tokens_to_ids(lowercase_) for lang_code in FAIRSEQ_LANGUAGE_CODES
          }

          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		src_lang if src_lang is not None else '''en_XX'''
          SCREAMING_SNAKE_CASE_      :  Any    						=		self.convert_tokens_to_ids(self._src_lang)
          SCREAMING_SNAKE_CASE_      :  int    						=		tgt_lang
          self.set_src_lang_special_tokens(self._src_lang)
      @property
      def _SCREAMING_SNAKE_CASE (						self     :  Any):




          '''simple docstring'''


          return self._src_lang
      @src_lang.setter
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple						,	lowercase_     :  str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Tuple    						=		new_src_lang
          self.set_src_lang_special_tokens(self._src_lang)
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]						,	lowercase_     :  List[int]						,	lowercase_     :  Optional[List[int]] = None):




          '''simple docstring'''


          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 _SCREAMING_SNAKE_CASE (						self     :  Dict						,	lowercase_     :  List[int]						,	lowercase_     :  Optional[List[int]] = None):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		[self.sep_token_id]
          SCREAMING_SNAKE_CASE_      :  Any    						=		[self.cls_token_id]

          if token_ids_a is None:
              return len(cls + token_ids_a + sep) * [0]
          return len(cls + token_ids_a + sep + sep + token_ids_a + sep) * [0]
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[Any]						,	lowercase_     :  Dict						,	lowercase_     :  str						,	lowercase_     :  Optional[str]						,	lowercase_     :  Optional[str]						,	**lowercase_     :  List[Any]):




          '''simple docstring'''


          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_      :  List[str]    						=		src_lang
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		self(lowercase_						,	add_special_tokens=lowercase_						,	return_tensors=lowercase_						,	**lowercase_)
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		self.convert_tokens_to_ids(lowercase_)
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		tgt_lang_id
          return inputs
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[Any]						,	lowercase_     :  List[str]						,	lowercase_     :  str = "en_XX"						,	lowercase_     :  Optional[List[str]] = None						,	lowercase_     :  str = "ro_RO"						,	**lowercase_     :  Any						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		src_lang
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		tgt_lang
          return super().prepare_seqaseq_batch(lowercase_						,	lowercase_						,	**lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]):




          '''simple docstring'''


          return self.set_src_lang_special_tokens(self.src_lang)
      def _SCREAMING_SNAKE_CASE (						self     :  Union[str, Any]):




          '''simple docstring'''


          return self.set_tgt_lang_special_tokens(self.tgt_lang)
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple						,	lowercase_     :  Union[str, Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  int    						=		self.convert_tokens_to_ids(lowercase_)
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		[]
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		[self.eos_token_id, self.cur_lang_code]

          SCREAMING_SNAKE_CASE_      :  List[str]    						=		self.convert_ids_to_tokens(self.prefix_tokens)
          SCREAMING_SNAKE_CASE_      :  str    						=		self.convert_ids_to_tokens(self.suffix_tokens)

          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		processors.TemplateProcessing(
              single=prefix_tokens_str + ['''$A'''] + suffix_tokens_str						,	pair=prefix_tokens_str + ['''$A''', '''$B'''] + suffix_tokens_str						,	special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str						,	self.prefix_tokens + self.suffix_tokens))						,	)
      def _SCREAMING_SNAKE_CASE (						self     :  Union[str, Any]						,	lowercase_     :  str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[str]    						=		self.convert_tokens_to_ids(lowercase_)
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		[]
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		[self.eos_token_id, self.cur_lang_code]

          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		self.convert_ids_to_tokens(self.prefix_tokens)
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		self.convert_ids_to_tokens(self.suffix_tokens)

          SCREAMING_SNAKE_CASE_      :  Dict    						=		processors.TemplateProcessing(
              single=prefix_tokens_str + ['''$A'''] + suffix_tokens_str						,	pair=prefix_tokens_str + ['''$A''', '''$B'''] + suffix_tokens_str						,	special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str						,	self.prefix_tokens + self.suffix_tokens))						,	)



      def _SCREAMING_SNAKE_CASE (						self     :  Dict						,	lowercase_     :  str						,	lowercase_     :  Optional[str] = None):




          '''simple docstring'''


          if not self.can_save_slow_tokenizer:
              raise ValueError(
                  '''Your fast tokenizer does not have the necessary information to save the vocabulary for a slow '''
                  '''tokenizer.''')

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

          if os.path.abspath(self.vocab_file) != os.path.abspath(lowercase_):
              copyfile(self.vocab_file						,	lowercase_)

          return (out_vocab_file,)










"""simple docstring"""






import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible




class 							lowerCAmelCase__				(			unittest.TestCase ):






      '''simple docstring'''

      def _SCREAMING_SNAKE_CASE (						self     :  str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		[
              '''safety_checker/pytorch_model.bin''',
              '''safety_checker/model.safetensors''',
              '''vae/diffusion_pytorch_model.bin''',
              '''vae/diffusion_pytorch_model.safetensors''',
              '''text_encoder/pytorch_model.bin''',
              '''text_encoder/model.safetensors''',
              '''unet/diffusion_pytorch_model.bin''',
              '''unet/diffusion_pytorch_model.safetensors''',
          ]
          self.assertTrue(is_safetensors_compatible(lowercase_))
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		[
              '''unet/diffusion_pytorch_model.bin''',
              '''unet/diffusion_pytorch_model.safetensors''',
          ]
          self.assertTrue(is_safetensors_compatible(lowercase_))
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		[
              '''safety_checker/pytorch_model.bin''',
              '''safety_checker/model.safetensors''',
              '''vae/diffusion_pytorch_model.bin''',
              '''vae/diffusion_pytorch_model.safetensors''',
              '''text_encoder/pytorch_model.bin''',
              '''text_encoder/model.safetensors''',
              '''unet/diffusion_pytorch_model.bin''',
              # Removed: 'unet/diffusion_pytorch_model.safetensors',
          ]
          self.assertFalse(is_safetensors_compatible(lowercase_))
      def _SCREAMING_SNAKE_CASE (						self     :  int):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		[
              '''text_encoder/pytorch_model.bin''',
              '''text_encoder/model.safetensors''',
          ]
          self.assertTrue(is_safetensors_compatible(lowercase_))
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		[
              '''safety_checker/pytorch_model.bin''',
              '''safety_checker/model.safetensors''',
              '''vae/diffusion_pytorch_model.bin''',
              '''vae/diffusion_pytorch_model.safetensors''',
              '''text_encoder/pytorch_model.bin''',
              # Removed: 'text_encoder/model.safetensors',
              '''unet/diffusion_pytorch_model.bin''',
              '''unet/diffusion_pytorch_model.safetensors''',
          ]
          self.assertFalse(is_safetensors_compatible(lowercase_))
      def _SCREAMING_SNAKE_CASE (						self     :  List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  str    						=		[
              '''safety_checker/pytorch_model.fp16.bin''',
              '''safety_checker/model.fp16.safetensors''',
              '''vae/diffusion_pytorch_model.fp16.bin''',
              '''vae/diffusion_pytorch_model.fp16.safetensors''',
              '''text_encoder/pytorch_model.fp16.bin''',
              '''text_encoder/model.fp16.safetensors''',
              '''unet/diffusion_pytorch_model.fp16.bin''',
              '''unet/diffusion_pytorch_model.fp16.safetensors''',
          ]
          SCREAMING_SNAKE_CASE_      :  Any    						=		'''fp16'''
          self.assertTrue(is_safetensors_compatible(lowercase_						,	variant=lowercase_))
      def _SCREAMING_SNAKE_CASE (						self     :  str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		[
              '''unet/diffusion_pytorch_model.fp16.bin''',
              '''unet/diffusion_pytorch_model.fp16.safetensors''',
          ]
          SCREAMING_SNAKE_CASE_      :  Dict    						=		'''fp16'''
          self.assertTrue(is_safetensors_compatible(lowercase_						,	variant=lowercase_))
      def _SCREAMING_SNAKE_CASE (						self     :  int):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[str]    						=		[
              '''unet/diffusion_pytorch_model.bin''',
              '''unet/diffusion_pytorch_model.safetensors''',
          ]
          SCREAMING_SNAKE_CASE_      :  Any    						=		'''fp16'''
          self.assertTrue(is_safetensors_compatible(lowercase_						,	variant=lowercase_))
      def _SCREAMING_SNAKE_CASE (						self     :  Union[str, Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Tuple    						=		[
              '''safety_checker/pytorch_model.fp16.bin''',
              '''safety_checker/model.fp16.safetensors''',
              '''vae/diffusion_pytorch_model.fp16.bin''',
              '''vae/diffusion_pytorch_model.fp16.safetensors''',
              '''text_encoder/pytorch_model.fp16.bin''',
              '''text_encoder/model.fp16.safetensors''',
              '''unet/diffusion_pytorch_model.fp16.bin''',
              # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
          ]
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		'''fp16'''
          self.assertFalse(is_safetensors_compatible(lowercase_						,	variant=lowercase_))
      def _SCREAMING_SNAKE_CASE (						self     :  List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		[
              '''text_encoder/pytorch_model.fp16.bin''',
              '''text_encoder/model.fp16.safetensors''',
          ]
          SCREAMING_SNAKE_CASE_      :  Any    						=		'''fp16'''
          self.assertTrue(is_safetensors_compatible(lowercase_						,	variant=lowercase_))
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		[
              '''text_encoder/pytorch_model.bin''',
              '''text_encoder/model.safetensors''',
          ]
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		'''fp16'''
          self.assertTrue(is_safetensors_compatible(lowercase_						,	variant=lowercase_))



      def _SCREAMING_SNAKE_CASE (						self     :  Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		[
              '''safety_checker/pytorch_model.fp16.bin''',
              '''safety_checker/model.fp16.safetensors''',
              '''vae/diffusion_pytorch_model.fp16.bin''',
              '''vae/diffusion_pytorch_model.fp16.safetensors''',
              '''text_encoder/pytorch_model.fp16.bin''',
              # 'text_encoder/model.fp16.safetensors',
              '''unet/diffusion_pytorch_model.fp16.bin''',
              '''unet/diffusion_pytorch_model.fp16.safetensors''',
          ]
          SCREAMING_SNAKE_CASE_      :  str    						=		'''fp16'''
          self.assertFalse(is_safetensors_compatible(lowercase_						,	variant=lowercase_))


