



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


lowerCamelCase_							:Dict       = logging.get_logger(__name__)

lowerCamelCase_							:Optional[int]       = '''▁'''

lowerCamelCase_							:str       = {'''vocab_file''': '''sentencepiece.bpe.model'''}

lowerCamelCase_							: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'''
        ),
    }
}

lowerCamelCase_							:Optional[int]       = {
    '''facebook/mbart-large-en-ro''': 1_0_2_4,
    '''facebook/mbart-large-cc25''': 1_0_2_4,
}

# fmt: off
lowerCamelCase_							:List[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							(      __UpperCAmelCase			):
	__SCREAMING_SNAKE_CASE       :		Union[str, Any]		      = VOCAB_FILES_NAMES
	__SCREAMING_SNAKE_CASE       :		Optional[int]		      = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	__SCREAMING_SNAKE_CASE       :		List[str]		      = PRETRAINED_VOCAB_FILES_MAP
	__SCREAMING_SNAKE_CASE       :		int		      = ['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	,				):
				# Mask token behave like a normal word, i.e. include the space before it
				A_     :						str     =					AddedToken(lowercase	,				lstrip=lowercase	,				rstrip=lowercase							) if isinstance(lowercase	,				lowercase							) else mask_token

				A_     :						int     =					{} 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	,				)

				A_     :						Optional[Any]     =					spm.SentencePieceProcessor(**self.sp_model_kwargs							)
				self.sp_model.Load(str(lowercase							)							)
				A_     :						Any     =					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
				A_     :						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
				A_     :						Dict     =					1

				A_     :						Dict     =					len(self.sp_model							)
				A_     :						str     =					{
				    code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(lowercase							)
				}
				A_     :						List[Any]     =					{v: k for k, v in self.lang_code_to_id.items()}
				A_     :						List[Any]     =					len(self.sp_model							) + len(self.lang_code_to_id							) + self.fairseq_offset

				self.fairseq_tokens_to_ids.update(self.lang_code_to_id							)
				A_     :						Dict     =					{v: k for k, v in self.fairseq_tokens_to_ids.items()}
				A_     :						Optional[int]     =					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]							)

				A_     :						Tuple     =					src_lang if src_lang is not None else """en_XX"""
				A_     :						List[str]     =					self.lang_code_to_id[self._src_lang]
				A_     :						int     =					tgt_lang
				self.set_src_lang_special_tokens(self._src_lang							)


	def __getstate__(self							):
				A_     :						Union[str, Any]     =					self.__dict__.copy()
				A_     :						Optional[Any]     =					None
				A_     :						Dict     =					self.sp_model.serialized_model_proto()
				return state


	def __setstate__(self	,				lowercase							):
				A_     :						Optional[int]     =					d

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

				A_     :						List[str]     =					spm.SentencePieceProcessor(**self.sp_model_kwargs							)
				self.sp_model.LoadFromSerializedProto(self.sp_model_proto							)


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


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


	@src_lang.setter
	def    _a     (self	,				lowercase							):
				A_     :						List[Any]     =					new_src_lang
				self.set_src_lang_special_tokens(self._src_lang							)


	def    _a     (self	,				lowercase	,				lowercase = None	,				lowercase = False							):

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

				A_     :						List[Any]     =					[1] * len(self.prefix_tokens							)
				A_     :						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    _a     (self	,				lowercase	,				lowercase = 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    _a     (self	,				lowercase	,				lowercase = None							):
				A_     :						Dict     =					[self.sep_token_id]
				A_     :						List[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    _a     (self	,				lowercase	,				lowercase	,				lowercase	,				lowercase	,				**lowercase							):
				if src_lang is None or tgt_lang is None:
							raise ValueError("""Translation requires a `src_lang` and a `tgt_lang` for this model"""							)
				A_     :						Any     =					src_lang
				A_     :						Optional[int]     =					self(lowercase	,				add_special_tokens=lowercase	,				return_tensors=lowercase	,				**lowercase							)
				A_     :						int     =					self.convert_tokens_to_ids(lowercase							)
				A_     :						Optional[int]     =					tgt_lang_id
				return inputs


	def    _a     (self							):
				A_     :						List[Any]     =					{self.convert_ids_to_tokens(lowercase							): i for i in range(self.vocab_size							)}
				vocab.update(self.added_tokens_encoder							)
				return vocab


	def    _a     (self	,				lowercase							):
				return self.sp_model.encode(lowercase	,				out_type=lowercase							)


	def    _a     (self	,				lowercase							):
				if token in self.fairseq_tokens_to_ids:
							return self.fairseq_tokens_to_ids[token]
				A_     :						Dict     =					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    _a     (self	,				lowercase							):
				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    _a     (self	,				lowercase							):
				A_     :						Any     =					"""""".join(lowercase							).replace(lowercase	,				""" """							).strip()
				return out_string


	def    _a     (self	,				lowercase	,				lowercase = None							):
				if not os.path.isdir(lowercase							):
							logger.error(F'Vocabulary path ({save_directory}) should be a directory'							)
							return
				A_     :						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:
										A_     :						int     =					self.sp_model.serialized_model_proto()
										fi.write(lowercase							)

				return (out_vocab_file,)


	def    _a     (self	,				lowercase	,				lowercase = "en_XX"	,				lowercase = None	,				lowercase = "ro_RO"	,				**lowercase	,				):
				A_     :						Any     =					src_lang
				A_     :						Tuple     =					tgt_lang
				return super().prepare_seqaseq_batch(lowercase	,				lowercase	,				**lowercase							)


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


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


	def    _a     (self	,				lowercase							):
				A_     :						str     =					self.lang_code_to_id[src_lang]
				A_     :						str     =					[]
				A_     :						Dict     =					[self.eos_token_id, self.cur_lang_code]


	def    _a     (self	,				lowercase							):
				A_     :						Optional[Any]     =					self.lang_code_to_id[lang]
				A_     :						int     =					[]
				A_     :						Union[str, Any]     =					[self.eos_token_id, self.cur_lang_code]




'''simple docstring'''




from typing import Dict

from .base import GenericTensor, Pipeline







class      _lowerCAmelCase							(      __UpperCAmelCase			):


	def    _a     (self	,				lowercase=None	,				lowercase=None	,				lowercase=None	,				**lowercase							):
				if tokenize_kwargs is None:
							A_     :						Optional[Any]     =					{}

				if truncation is not None:
							if "truncation" in tokenize_kwargs:
										raise ValueError(
										    """truncation parameter defined twice (given as keyword argument as well as in tokenize_kwargs)"""							)
							A_     :						str     =					truncation

				A_     :						List[str]     =					tokenize_kwargs

				A_     :						Dict     =					{}
				if return_tensors is not None:
							A_     :						List[Any]     =					return_tensors

				return preprocess_params, {}, postprocess_params


	def    _a     (self	,				lowercase	,				**lowercase							):
				A_     :						Optional[int]     =					self.framework
				A_     :						str     =					self.tokenizer(lowercase	,				return_tensors=lowercase	,				**lowercase							)
				return model_inputs


	def    _a     (self	,				lowercase							):
				A_     :						str     =					self.model(**lowercase							)
				return model_outputs


	def    _a     (self	,				lowercase	,				lowercase=False							):
				# [0] is the first available tensor, logits or last_hidden_state.
				if return_tensors:
							return model_outputs[0]
				if self.framework == "pt":
							return model_outputs[0].tolist()
				elif self.framework == "tf":
							return model_outputs[0].numpy().tolist()


	def __call__(self	,				*lowercase	,				**lowercase							):
				return super().__call__(*lowercase	,				**lowercase							)