


"""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


UpperCAmelCase    :     str   						=		logging.get_logger(__name__)

UpperCAmelCase    :     Union[str, Any]   						=		'▁'

UpperCAmelCase    :     Dict   						=		{'vocab_file': 'sentencepiece.bpe.model'}

UpperCAmelCase    :     str   						=		{
    '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'
        ),
    }
}

UpperCAmelCase    :     str   						=		{
    'facebook/mbart-large-en-ro': 1024,
    'facebook/mbart-large-cc25': 1024,
}

# fmt: off
UpperCAmelCase    :     str   						=		['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__   (     A				):







					"""simple docstring"""


					__a	       =							VOCAB_FILES_NAMES
					__a	       =							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					__a	       =							PRETRAINED_VOCAB_FILES_MAP
					__a	       =							["""input_ids""", """attention_mask"""]

					__a	       =							[]
					__a	       =							[]






					def __init__(       self    : Optional[Any] ,       UpperCamelCase    : Optional[int] ,       UpperCamelCase    : int="<s>" ,       UpperCamelCase    : int="</s>" ,       UpperCamelCase    : Optional[int]="</s>" ,       UpperCamelCase    : Any="<s>" ,       UpperCamelCase    : int="<unk>" ,       UpperCamelCase    : Dict="<pad>" ,       UpperCamelCase    : List[str]="<mask>" ,       UpperCamelCase    : Dict=None ,       UpperCamelCase    : int=None ,       UpperCamelCase    : Union[str, Any]=None ,       UpperCamelCase    : Optional[Dict[str, Any]] = None ,       UpperCamelCase    : List[Any]=None ,       **UpperCamelCase    : Optional[int] ,       ):
												'''simple docstring'''






												__UpperCAmelCase   :						int      							=					AddedToken(UpperCamelCase ,       lstrip=UpperCamelCase ,       rstrip=UpperCamelCase  ) if isinstance(UpperCamelCase ,       UpperCamelCase  ) else mask_token

												__UpperCAmelCase   :						Any      							=					{} if sp_model_kwargs is None else sp_model_kwargs

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

												__UpperCAmelCase   :						List[Any]      							=					spm.SentencePieceProcessor(**self.sp_model_kwargs  )
												self.sp_model.Load(str(UpperCamelCase  )  )
												__UpperCAmelCase   :						Dict      							=					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   :						Optional[Any]      							=					{"""<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   :						int      							=					1

												__UpperCAmelCase   :						Any      							=					len(self.sp_model  )
												__UpperCAmelCase   :						int      							=					{
												    code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(UpperCamelCase  )
												}
												__UpperCAmelCase   :						Dict      							=					{v: k for k, v in self.lang_code_to_id.items()}
												__UpperCAmelCase   :						Tuple      							=					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   :						int      							=					{v: k for k, v in self.fairseq_tokens_to_ids.items()}
												__UpperCAmelCase   :						Union[str, Any]      							=					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   :						Union[str, Any]      							=					src_lang if src_lang is not None else """en_XX"""
												__UpperCAmelCase   :						Any      							=					self.lang_code_to_id[self._src_lang]
												__UpperCAmelCase   :						str      							=					tgt_lang
												self.set_src_lang_special_tokens(self._src_lang  )






					def __getstate__(       self    : int  ):
												'''simple docstring'''






												__UpperCAmelCase   :						int      							=					self.__dict__.copy()
												__UpperCAmelCase   :						Dict      							=					None
												__UpperCAmelCase   :						Any      							=					self.sp_model.serialized_model_proto()
												return state






					def __setstate__(       self    : Optional[Any] ,       UpperCamelCase    : int  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Union[str, Any]      							=					d

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

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






					@property
					def lowerCamelCase__							(       self    : List[str]  ):
												'''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 lowerCamelCase__							(       self    : Optional[int]  ):
												'''simple docstring'''






												return self._src_lang






					@src_lang.setter
					def lowerCamelCase__							(       self    : List[Any] ,       UpperCamelCase    : str  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Any      							=					new_src_lang
												self.set_src_lang_special_tokens(self._src_lang  )






					def lowerCamelCase__							(       self    : str ,       UpperCamelCase    : List[int] ,       UpperCamelCase    : Optional[List[int]] = None ,       UpperCamelCase    : bool = False  ):
												'''simple docstring'''






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

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






					def lowerCamelCase__							(       self    : List[str] ,       UpperCamelCase    : List[int] ,       UpperCamelCase    : 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 lowerCamelCase__							(       self    : Any ,       UpperCamelCase    : List[int] ,       UpperCamelCase    : Optional[List[int]] = None  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Optional[Any]      							=					[self.sep_token_id]
												__UpperCAmelCase   :						Optional[int]      							=					[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 lowerCamelCase__							(       self    : Any ,       UpperCamelCase    : Any ,       UpperCamelCase    : str ,       UpperCamelCase    : Optional[str] ,       UpperCamelCase    : Optional[str] ,       **UpperCamelCase    : Dict  ):
												'''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   :						str      							=					src_lang
												__UpperCAmelCase   :						str      							=					self(UpperCamelCase ,       add_special_tokens=UpperCamelCase ,       return_tensors=UpperCamelCase ,       **UpperCamelCase  )
												__UpperCAmelCase   :						Dict      							=					self.convert_tokens_to_ids(UpperCamelCase  )
												__UpperCAmelCase   :						List[str]      							=					tgt_lang_id
												return inputs






					def lowerCamelCase__							(       self    : Any  ):
												'''simple docstring'''






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






					def lowerCamelCase__							(       self    : str ,       UpperCamelCase    : str  ):
												'''simple docstring'''






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






					def lowerCamelCase__							(       self    : Optional[Any] ,       UpperCamelCase    : Tuple  ):
												'''simple docstring'''






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

												# 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    : List[str] ,       UpperCamelCase    : Dict  ):
												'''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 lowerCamelCase__							(       self    : Optional[Any] ,       UpperCamelCase    : List[str]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						List[Any]      							=					"""""".join(UpperCamelCase  ).replace(UpperCamelCase ,       """ """  ).strip()
												return out_string






					def lowerCamelCase__							(       self    : Any ,       UpperCamelCase    : str ,       UpperCamelCase    : Optional[str] = None  ):
												'''simple docstring'''






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

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

												return (out_vocab_file,)






					def lowerCamelCase__							(       self    : int ,       UpperCamelCase    : List[str] ,       UpperCamelCase    : str = "en_XX" ,       UpperCamelCase    : Optional[List[str]] = None ,       UpperCamelCase    : str = "ro_RO" ,       **UpperCamelCase    : Optional[Any] ,       ):
												'''simple docstring'''






												__UpperCAmelCase   :						List[Any]      							=					src_lang
												__UpperCAmelCase   :						List[str]      							=					tgt_lang
												return super().prepare_seqaseq_batch(UpperCamelCase ,       UpperCamelCase ,       **UpperCamelCase  )






					def lowerCamelCase__							(       self    : int  ):
												'''simple docstring'''






												return self.set_src_lang_special_tokens(self.src_lang  )






					def lowerCamelCase__							(       self    : Tuple  ):
												'''simple docstring'''






												return self.set_tgt_lang_special_tokens(self.tgt_lang  )






					def lowerCamelCase__							(       self    : Optional[Any] ,       UpperCamelCase    : Union[str, Any]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Dict      							=					self.lang_code_to_id[src_lang]
												__UpperCAmelCase   :						int      							=					[]
												__UpperCAmelCase   :						Dict      							=					[self.eos_token_id, self.cur_lang_code]






					def lowerCamelCase__							(       self    : Optional[int] ,       UpperCamelCase    : str  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Any      							=					self.lang_code_to_id[lang]
												__UpperCAmelCase   :						Union[str, Any]      							=					[]
												__UpperCAmelCase   :						Optional[Any]      							=					[self.eos_token_id, self.cur_lang_code]




"""simple docstring"""






from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
							if not (is_transformers_available() and is_torch_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from ...utils.dummy_torch_and_transformers_objects import (
							    ImageTextPipelineOutput,
							    UniDiffuserPipeline,
							)
else:
							from .modeling_text_decoder import UniDiffuserTextDecoder
							from .modeling_uvit import UniDiffuserModel, UTransformeraDModel
							from .pipeline_unidiffuser import ImageTextPipelineOutput, UniDiffuserPipeline
