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

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


a_							=      logging.get_logger(__name__)

a_							=      """▁"""

a_							=      {"""vocab_file""": """sentencepiece.bpe.model"""}

a_							=      {
    """vocab_file""": {
        """facebook/xglm-564M""": """https://huggingface.co/facebook/xglm-564M/resolve/main/sentencepiece.bpe.model""",
    }
}

a_							=      {
    """facebook/xglm-564M""": 2048,
}


class 			UpperCAmelCase__    (		snake_case					):



  """simple docstring"""






  lowerCAmelCase__		:	str    		= VOCAB_FILES_NAMES
  lowerCAmelCase__		:	Optional[Any]    		= PRETRAINED_VOCAB_FILES_MAP
  lowerCAmelCase__		:	List[str]    		= PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  lowerCAmelCase__		:	Union[str, Any]    		= ['input_ids', 'attention_mask']






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




    '''simple docstring'''




    __UpperCAmelCase						    =       {} if sp_model_kwargs is None else sp_model_kwargs

    # Compatibility with the original tokenizer
    __UpperCAmelCase						    =       7
    __UpperCAmelCase						    =       [F'''<madeupword{i}>''' for i in range(self.num_madeup_words   )]

    __UpperCAmelCase						    =       kwargs.get("additional_special_tokens"				,				[]   )
    kwargs["additional_special_tokens"] += [
        word for word in madeup_words if word not in kwargs["additional_special_tokens"]
    ]

    super().__init__(
        bos_token=__lowerCAmelCase				,				eos_token=__lowerCAmelCase				,				unk_token=__lowerCAmelCase				,				sep_token=__lowerCAmelCase				,				cls_token=__lowerCAmelCase				,				pad_token=__lowerCAmelCase				,				sp_model_kwargs=self.sp_model_kwargs				,				**__lowerCAmelCase				,				)

    __UpperCAmelCase						    =       spm.SentencePieceProcessor(**self.sp_model_kwargs   )
    self.sp_model.Load(str(__lowerCAmelCase   )   )
    __UpperCAmelCase						    =       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'

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

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

    __UpperCAmelCase						    =       len(self.sp_model   )
    __UpperCAmelCase						    =       {F'''<madeupword{i}>''': sp_size + i + self.fairseq_offset for i in range(self.num_madeup_words   )}
    self.fairseq_tokens_to_ids.update(__lowerCAmelCase   )

    __UpperCAmelCase						    =       {v: k for k, v in self.fairseq_tokens_to_ids.items()}






  def __getstate__(				self:			Tuple   ) -> int:




    '''simple docstring'''




    __UpperCAmelCase						    =       self.__dict__.copy()
    __UpperCAmelCase						    =       None
    __UpperCAmelCase						    =       self.sp_model.serialized_model_proto()
    return state






  def __setstate__(				self:			str				,				__lowerCAmelCase:			int   ) -> Optional[Any]:




    '''simple docstring'''




    __UpperCAmelCase						    =       d

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

    __UpperCAmelCase						    =       spm.SentencePieceProcessor(**self.sp_model_kwargs   )
    self.sp_model.LoadFromSerializedProto(self.sp_model_proto   )






  def 	_UpperCAmelCase				(				self:			Optional[int]				,				__lowerCAmelCase:			List[int]				,				__lowerCAmelCase:			Optional[List[int]] = None   ) -> List[int]:




    '''simple docstring'''




    if token_ids_a is None:
      return [self.sep_token_id] + token_ids_a
    __UpperCAmelCase						    =       [self.sep_token_id]
    return sep + token_ids_a + sep + sep + token_ids_a






  def 	_UpperCAmelCase				(				self:			Union[str, Any]				,				__lowerCAmelCase:			List[int]				,				__lowerCAmelCase:			Optional[List[int]] = None				,				__lowerCAmelCase:			bool = False   ) -> List[int]:




    '''simple docstring'''




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

    if token_ids_a is None:
      return [1] + ([0] * len(__lowerCAmelCase   ))
    return [1] + ([0] * len(__lowerCAmelCase   )) + [1, 1] + ([0] * len(__lowerCAmelCase   ))






  def 	_UpperCAmelCase				(				self:			int				,				__lowerCAmelCase:			List[int]				,				__lowerCAmelCase:			Optional[List[int]] = None   ) -> List[int]:




    '''simple docstring'''




    __UpperCAmelCase						    =       [self.sep_token_id]

    if token_ids_a is None:
      return len(sep + token_ids_a   ) * [0]
    return len(sep + token_ids_a + sep + sep + token_ids_a   ) * [0]






  @property
  def 	_UpperCAmelCase				(				self:			List[Any]   ) -> str:




    '''simple docstring'''




    return len(self.sp_model   ) + self.fairseq_offset + self.num_madeup_words






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




    '''simple docstring'''




    __UpperCAmelCase						    =       {self.convert_ids_to_tokens(__lowerCAmelCase   ): i for i in range(self.vocab_size   )}
    vocab.update(self.added_tokens_encoder   )
    return vocab






  def 	_UpperCAmelCase				(				self:			Any				,				__lowerCAmelCase:			str   ) -> List[str]:




    '''simple docstring'''




    return self.sp_model.encode(__lowerCAmelCase				,				out_type=__lowerCAmelCase   )






  def 	_UpperCAmelCase				(				self:			List[Any]				,				__lowerCAmelCase:			Optional[Any]   ) -> List[str]:




    '''simple docstring'''




    if token in self.fairseq_tokens_to_ids:
      return self.fairseq_tokens_to_ids[token]
    __UpperCAmelCase						    =       self.sp_model.PieceToId(__lowerCAmelCase   )

    # 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 	_UpperCAmelCase				(				self:			Dict				,				__lowerCAmelCase:			str   ) -> Tuple:




    '''simple docstring'''




    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 	_UpperCAmelCase				(				self:			Union[str, Any]				,				__lowerCAmelCase:			Optional[Any]   ) -> List[Any]:




    '''simple docstring'''




    __UpperCAmelCase						    =       "".join(__lowerCAmelCase   ).replace(__lowerCAmelCase				,				" "   ).strip()
    return out_string






  def 	_UpperCAmelCase				(				self:			int				,				__lowerCAmelCase:			str				,				__lowerCAmelCase:			Optional[str] = None   ) -> Tuple[str]:




    '''simple docstring'''




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

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

    return (out_vocab_file,)




import inspect
import unittest


class 			UpperCAmelCase__    (		unittest.TestCase					):



  """simple docstring"""






  def 	_UpperCAmelCase				(				self:			Union[str, Any]   ) -> Dict:




    '''simple docstring'''




    try:
      import diffusers  # noqa: F401
    except ImportError:
      assert False






  def 	_UpperCAmelCase				(				self:			Dict   ) -> Dict:




    '''simple docstring'''




    import diffusers
    from diffusers.dependency_versions_table import deps

    __UpperCAmelCase						    =       inspect.getmembers(__lowerCAmelCase				,				inspect.isclass   )

    for cls_name, cls_module in all_classes:
      if "dummy_" in cls_module.__module__:
        for backend in cls_module._backends:
          if backend == "k_diffusion":
            __UpperCAmelCase						    =       "k-diffusion"
          elif backend == "invisible_watermark":
            __UpperCAmelCase						    =       "invisible-watermark"
          assert backend in deps, F'''{backend} is not in the deps table!'''



