
"""simple docstring"""





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


_a		       =							logging.get_logger(__name__)

_a		       =							'▁'

_a		       =							{'vocab_file': 'sentencepiece.bpe.model'}

_a		       =							{
    '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'
        ),
    }
}

_a		       =							{
    'facebook/mbart-large-en-ro': 1_024,
    'facebook/mbart-large-cc25': 1_024,
}

# fmt: off
_a		       =							['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 A_	(lowercase__      ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__			:						Dict       =		VOCAB_FILES_NAMES
  SCREAMING_SNAKE_CASE__			:						Optional[int]       =		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  SCREAMING_SNAKE_CASE__			:						Optional[Any]       =		PRETRAINED_VOCAB_FILES_MAP
  SCREAMING_SNAKE_CASE__			:						Optional[Any]       =		["""input_ids""", """attention_mask"""]

  SCREAMING_SNAKE_CASE__			:						List[int]       =		[]
  SCREAMING_SNAKE_CASE__			:						List[int]       =		[]







  def __init__(		self				,   lowercase_				,   lowercase_="<s>"				,   lowercase_="</s>"				,   lowercase_="</s>"				,   lowercase_="<s>"				,   lowercase_="<unk>"				,   lowercase_="<pad>"				,   lowercase_="<mask>"				,   lowercase_=None				,   lowercase_=None				,   lowercase_=None				,   lowercase_ = None				,   lowercase_=None				,   **lowercase_				,   ):

     """simple docstring"""
     # Mask token behave like a normal word, i.e. include the space before it
     UpperCAmelCase_					:					List[Any]			     =    AddedToken(lowercase_				,   lstrip=lowercase_				,   rstrip=lowercase_  ) if isinstance(lowercase_				,   lowercase_  ) else mask_token

     UpperCAmelCase_					:					Union[str, Any]			     =    {} if sp_model_kwargs is None else sp_model_kwargs

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

     UpperCAmelCase_					:					Dict			     =    spm.SentencePieceProcessor(**self.sp_model_kwargs  )
     self.sp_model.Load(str(lowercase_  )  )
     UpperCAmelCase_					:					str			     =    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
     UpperCAmelCase_					:					List[str]			     =    {"<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
     UpperCAmelCase_					:					Tuple			     =    1

     UpperCAmelCase_					:					Tuple			     =    len(self.sp_model  )
     UpperCAmelCase_					:					Tuple			     =    {
         code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(lowercase_  )
     }
     UpperCAmelCase_					:					int			     =    {v: k for k, v in self.lang_code_to_id.items()}
     UpperCAmelCase_					:					int			     =    len(self.sp_model  ) + len(self.lang_code_to_id  ) + self.fairseq_offset

     self.fairseq_tokens_to_ids.update(self.lang_code_to_id  )
     UpperCAmelCase_					:					Optional[Any]			     =    {v: k for k, v in self.fairseq_tokens_to_ids.items()}
     UpperCAmelCase_					:					str			     =    list(self.lang_code_to_id.keys()  )

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

     UpperCAmelCase_					:					str			     =    src_lang if src_lang is not None else "en_XX"
     UpperCAmelCase_					:					Any			     =    self.lang_code_to_id[self._src_lang]
     UpperCAmelCase_					:					Optional[Any]			     =    tgt_lang
     self.set_src_lang_special_tokens(self._src_lang  )







  def __getstate__(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[Any]			     =    self.__dict__.copy()
     UpperCAmelCase_					:					List[str]			     =    None
     UpperCAmelCase_					:					List[Any]			     =    self.sp_model.serialized_model_proto()
     return state







  def __setstate__(		self				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[Any]			     =    d

     # for backward compatibility
     if not hasattr(self				,   "sp_model_kwargs"  ):
        UpperCAmelCase_					:					List[str]			     =    {}

     UpperCAmelCase_					:					str			     =    spm.SentencePieceProcessor(**self.sp_model_kwargs  )
     self.sp_model.LoadFromSerializedProto(self.sp_model_proto  )







  @property
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     return len(self.sp_model  ) + len(self.lang_code_to_id  ) + self.fairseq_offset + 1  # Plus 1 for the mask token







  @property
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     return self._src_lang







  @src_lang.setter
  def        UpperCamelCase__					(		self				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					str			     =    new_src_lang
     self.set_src_lang_special_tokens(self._src_lang  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_ = None				,   lowercase_ = False  ):

     """simple docstring"""

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

     UpperCAmelCase_					:					Optional[Any]			     =    [1] * len(self.prefix_tokens  )
     UpperCAmelCase_					:					Dict			     =    [1] * len(self.suffix_tokens  )
     if token_ids_a is None:
        return prefix_ones + ([0] * len(lowercase_  )) + suffix_ones
     return prefix_ones + ([0] * len(lowercase_  )) + ([0] * len(lowercase_  )) + suffix_ones







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_ = 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        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_ = None  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[int]			     =    [self.sep_token_id]
     UpperCAmelCase_					:					Tuple			     =    [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        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_				,   **lowercase_  ):

     """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"  )
     UpperCAmelCase_					:					List[str]			     =    src_lang
     UpperCAmelCase_					:					Any			     =    self(lowercase_				,   add_special_tokens=lowercase_				,   return_tensors=lowercase_				,   **lowercase_  )
     UpperCAmelCase_					:					Union[str, Any]			     =    self.convert_tokens_to_ids(lowercase_  )
     UpperCAmelCase_					:					Union[str, Any]			     =    tgt_lang_id
     return inputs







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Union[str, Any]			     =    {self.convert_ids_to_tokens(lowercase_  ): i for i in range(self.vocab_size  )}
     vocab.update(self.added_tokens_encoder  )
     return vocab







  def        UpperCamelCase__					(		self				,   lowercase_  ):

     """simple docstring"""
     return self.sp_model.encode(lowercase_				,   out_type=lowercase_  )







  def        UpperCamelCase__					(		self				,   lowercase_  ):

     """simple docstring"""
     if token in self.fairseq_tokens_to_ids:
        return self.fairseq_tokens_to_ids[token]
     UpperCAmelCase_					:					Tuple			     =    self.sp_model.PieceToId(lowercase_  )

     # 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				,   lowercase_  ):

     """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				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					str			     =    "".join(lowercase_  ).replace(lowercase_				,   " "  ).strip()
     return out_string







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_ = None  ):

     """simple docstring"""
     if not os.path.isdir(lowercase_  ):
        logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""  )
        return
     UpperCAmelCase_					:					Union[str, Any]			     =    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_  ) and os.path.isfile(self.vocab_file  ):
        copyfile(self.vocab_file				,   lowercase_  )
     elif not os.path.isfile(self.vocab_file  ):
        with open(lowercase_				,   "wb"  ) as fi:
           UpperCAmelCase_					:					List[Any]			     =    self.sp_model.serialized_model_proto()
           fi.write(lowercase_  )

     return (out_vocab_file,)







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_ = "en_XX"				,   lowercase_ = None				,   lowercase_ = "ro_RO"				,   **lowercase_				,   ):

     """simple docstring"""
     UpperCAmelCase_					:					Union[str, Any]			     =    src_lang
     UpperCAmelCase_					:					Union[str, Any]			     =    tgt_lang
     return super().prepare_seqaseq_batch(lowercase_				,   lowercase_				,   **lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     return self.set_src_lang_special_tokens(self.src_lang  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     return self.set_tgt_lang_special_tokens(self.tgt_lang  )







  def        UpperCamelCase__					(		self				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[str]			     =    self.lang_code_to_id[src_lang]
     UpperCAmelCase_					:					List[str]			     =    []
     UpperCAmelCase_					:					Optional[int]			     =    [self.eos_token_id, self.cur_lang_code]







  def        UpperCamelCase__					(		self				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[str]			     =    self.lang_code_to_id[lang]
     UpperCAmelCase_					:					Tuple			     =    []
     UpperCAmelCase_					:					Optional[Any]			     =    [self.eos_token_id, self.cur_lang_code]









from queue import Queue
from typing import TYPE_CHECKING, Optional


if TYPE_CHECKING:
     from ..models.auto import AutoTokenizer







class 				__lowerCAmelCase       :







      def    _lowercase  (    self		,      lowerCAmelCase__							)   ->						Optional[Any]:



        '''simple docstring'''




        raise NotImplementedError()







      def    _lowercase  (    self							)   ->						int:



        '''simple docstring'''




        raise NotImplementedError()







class 				__lowerCAmelCase       (    UpperCamelCase__):







      def __init__(    self		,      lowerCAmelCase__		,      lowerCAmelCase__ = False		,      **lowerCAmelCase__							)   ->						int:



        '''simple docstring'''




        a__	:					str										=tokenizer
        a__	:					List[str]										=skip_prompt
        a__	:					List[Any]										=decode_kwargs

        # variables used in the streaming process
        a__	:					Dict										=[]
        a__	:					int										=0
        a__	:					str										=True







      def    _lowercase  (    self		,      lowerCAmelCase__							)   ->						Union[str, Any]:



        '''simple docstring'''




        if len(value.shape							) > 1 and value.shape[0] > 1:
          raise ValueError("TextStreamer only supports batch size 1"							)
        elif len(value.shape							) > 1:
          a__	:					Any										=value[0]

        if self.skip_prompt and self.next_tokens_are_prompt:
          a__	:					Dict										=False
          return

        # Add the new token to the cache and decodes the entire thing.
        self.token_cache.extend(value.tolist()							)
        a__	:					Union[str, Any]										=self.tokenizer.decode(self.token_cache		,      **self.decode_kwargs							)

        # After the symbol for a new line, we flush the cache.
        if text.endswith("\n"							):
          a__	:					List[Any]										=text[self.print_len :]
          a__	:					List[str]										=[]
          a__	:					Optional[int]										=0
        # If the last token is a CJK character, we print the characters.
        elif len(lowerCAmelCase__							) > 0 and self._is_chinese_char(ord(text[-1]							)							):
          a__	:					List[str]										=text[self.print_len :]
          self.print_len += len(lowerCAmelCase__							)
        # Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words,
        # which may change with the subsequent token -- there are probably smarter ways to do this!)
        else:
          a__	:					str										=text[self.print_len : text.rfind(" "							) + 1]
          self.print_len += len(lowerCAmelCase__							)

        self.on_finalized_text(lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						Any:



        '''simple docstring'''




        if len(self.token_cache							) > 0:
          a__	:					Union[str, Any]										=self.tokenizer.decode(self.token_cache		,      **self.decode_kwargs							)
          a__	:					List[Any]										=text[self.print_len :]
          a__	:					List[str]										=[]
          a__	:					Optional[int]										=0
        else:
          a__	:					Union[str, Any]										=""

        a__	:					Any										=True
        self.on_finalized_text(lowerCAmelCase__		,      stream_end=lowerCAmelCase__							)







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__ = False							)   ->						Optional[Any]:



        '''simple docstring'''




        print(lowerCAmelCase__		,      flush=lowerCAmelCase__		,      end="" if not stream_end else None							)







      def    _lowercase  (    self		,      lowerCAmelCase__							)   ->						str:



        '''simple docstring'''




        if (
            (cp >= 0X4E_00 and cp <= 0X9F_FF)
            or (cp >= 0X34_00 and cp <= 0X4D_BF)  #
            or (cp >= 0X2_00_00 and cp <= 0X2_A6_DF)  #
            or (cp >= 0X2_A7_00 and cp <= 0X2_B7_3F)  #
            or (cp >= 0X2_B7_40 and cp <= 0X2_B8_1F)  #
            or (cp >= 0X2_B8_20 and cp <= 0X2_CE_AF)  #
            or (cp >= 0XF9_00 and cp <= 0XFA_FF)
            or (cp >= 0X2_F8_00 and cp <= 0X2_FA_1F)  #
        ):  #
          return True

        return False







class 				__lowerCAmelCase       (    UpperCamelCase__):







      def __init__(    self		,      lowerCAmelCase__		,      lowerCAmelCase__ = False		,      lowerCAmelCase__ = None		,      **lowerCAmelCase__							)   ->						Optional[Any]:



        '''simple docstring'''




        super().__init__(lowerCAmelCase__		,      lowerCAmelCase__		,      **lowerCAmelCase__							)
        a__	:					str										=Queue()
        a__	:					Optional[Any]										=None
        a__	:					Any										=timeout







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__ = False							)   ->						List[str]:



        '''simple docstring'''




        self.text_queue.put(lowerCAmelCase__		,      timeout=self.timeout							)
        if stream_end:
          self.text_queue.put(self.stop_signal		,      timeout=self.timeout							)







      def __iter__(    self							)   ->						Dict:



        '''simple docstring'''




        return self







      def    _lowercase  (    self							)   ->						int:



        '''simple docstring'''




        a__	:					int										=self.text_queue.get(timeout=self.timeout							)
        if value == self.stop_signal:
          raise StopIteration()
        else:
          return value






