





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					(					lowerCAmelCase				):

			'''simple docstring'''







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

			lowerCAmelCase_							  =		[]
			lowerCAmelCase_							  =		[]





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



				'''simple docstring'''







				lowerCamelCase__              =	AddedToken(__lowerCAmelCase ,						lstrip=__lowerCAmelCase ,						rstrip=__lowerCAmelCase							) if isinstance(__lowerCAmelCase ,						__lowerCAmelCase							) else mask_token

				lowerCamelCase__              =	{} if sp_model_kwargs is None else sp_model_kwargs

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

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

				lowerCamelCase__              =	len(self.sp_model							)
				lowerCamelCase__              =	{
				    code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(__lowerCAmelCase							)
				}
				lowerCamelCase__              =	{v: k for k, v in self.lang_code_to_id.items()}
				lowerCamelCase__              =	len(self.sp_model							) + len(self.lang_code_to_id							) + self.fairseq_offset

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

				lowerCamelCase__              =	src_lang if src_lang is not None else '''en_XX'''
				lowerCamelCase__              =	self.lang_code_to_id[self._src_lang]
				lowerCamelCase__              =	tgt_lang
				self.set_src_lang_special_tokens(self._src_lang							)





			def __getstate__(					self							):



				'''simple docstring'''







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





			def __setstate__(					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	d

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

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





			@property
			def 	__lowerCamelCase       (					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 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				return self._src_lang





			@src_lang.setter
			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	new_src_lang
				self.set_src_lang_special_tokens(self._src_lang							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase = None ,						__lowerCAmelCase = False							):



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

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





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



				'''simple docstring'''







				lowerCamelCase__              =	[self.sep_token_id]
				lowerCamelCase__              =	[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 ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						**__lowerCAmelCase							):



				'''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'''							)
				lowerCamelCase__              =	src_lang
				lowerCamelCase__              =	self(__lowerCAmelCase ,						add_special_tokens=__lowerCAmelCase ,						return_tensors=__lowerCAmelCase ,						**__lowerCAmelCase							)
				lowerCamelCase__              =	self.convert_tokens_to_ids(__lowerCAmelCase							)
				lowerCamelCase__              =	tgt_lang_id
				return inputs





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







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





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







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





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				if token in self.fairseq_tokens_to_ids:
					return self.fairseq_tokens_to_ids[token]
				lowerCamelCase__              =	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 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



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



				'''simple docstring'''







				lowerCamelCase__              =	''''''.join(__lowerCAmelCase							).replace(__lowerCAmelCase ,						''' '''							).strip()
				return out_string





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase = None							):



				'''simple docstring'''







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

				return (out_vocab_file,)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase = "en_XX" ,						__lowerCAmelCase = None ,						__lowerCAmelCase = "ro_RO" ,						**__lowerCAmelCase ,						):



				'''simple docstring'''







				lowerCamelCase__              =	src_lang
				lowerCamelCase__              =	tgt_lang
				return super().prepare_seqaseq_batch(__lowerCAmelCase ,						__lowerCAmelCase ,						**__lowerCAmelCase							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				return self.set_src_lang_special_tokens(self.src_lang							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				return self.set_tgt_lang_special_tokens(self.tgt_lang							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	self.lang_code_to_id[src_lang]
				lowerCamelCase__              =	[]
				lowerCamelCase__              =	[self.eos_token_id, self.cur_lang_code]





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	self.lang_code_to_id[lang]
				lowerCamelCase__              =	[]
				lowerCamelCase__              =	[self.eos_token_id, self.cur_lang_code]









import contextlib
from multiprocessing import Pool, RLock

from tqdm.auto import tqdm

from ..utils import experimental, logging


_a						=   logging.get_logger(__name__)






class       __A					:

			'''simple docstring'''







			lowerCAmelCase_							  =		None





@experimental
def lowerCAmelCase__(__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case							)       ->   Tuple:







	'''simple docstring'''







	if ParallelBackendConfig.backend_name is None:
		return _map_with_multiprocessing_pool(
		    __snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case							)

	return _map_with_joblib(__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case							)





def lowerCAmelCase__(__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case							)       ->   int:







	'''simple docstring'''







	lowerCamelCase__              =	num_proc if num_proc <= len(__snake_case							) else len(__snake_case							)
	lowerCamelCase__              =	[]  # We organize the splits ourselve (contiguous splits)
	for index in range(__snake_case							):
		lowerCamelCase__              =	len(__snake_case							) // num_proc
		lowerCamelCase__              =	len(__snake_case							) % num_proc
		lowerCamelCase__              =	div * index + min(__snake_case ,__snake_case							)
		lowerCamelCase__              =	start + div + (1 if index < mod else 0)
		split_kwds.append((function, iterable[start:end], types, index, disable_tqdm, desc)							)

	if len(__snake_case							) != sum(len(i[1]							) for i in split_kwds							):
		raise ValueError(
		    F'Error dividing inputs iterable among processes. '
		    F'Total number of objects {len(__snake_case							)}, '
		    F'length: {sum(len(i[1]							) for i in split_kwds							)}'							)

	logger.info(
	    F'Spawning {num_proc} processes for {len(__snake_case							)} objects in slices of {[len(i[1]							) for i in split_kwds]}'							)
	lowerCamelCase__       ,			lowerCamelCase__              =	None, None
	if not disable_tqdm:
		lowerCamelCase__       ,			lowerCamelCase__              =	(RLock(),), tqdm.set_lock
	with Pool(__snake_case ,initargs=__snake_case ,initializer=__snake_case							) as pool:
		lowerCamelCase__              =	pool.map(__snake_case ,__snake_case							)
	logger.info(F'Finished {num_proc} processes'							)
	lowerCamelCase__              =	[obj for proc_res in mapped for obj in proc_res]
	logger.info(F'Unpacked {len(__snake_case							)} objects'							)

	return mapped





def lowerCAmelCase__(__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case							)       ->   List[str]:







	'''simple docstring'''







	import joblib

	with joblib.parallel_backend(ParallelBackendConfig.backend_name ,n_jobs=__snake_case							):
		return joblib.Parallel()(
		    joblib.delayed(__snake_case							)((function, obj, types, None, True, None)							) for obj in iterable							)





@experimental
@contextlib.contextmanager
def lowerCAmelCase__(__snake_case							)       ->   int:







	'''simple docstring'''







	lowerCamelCase__              =	backend_name

	if backend_name == "spark":
		from joblibspark import register_spark

		register_spark()

		# TODO: call create_cache_and_write_probe if "download" in steps
		# TODO: raise NotImplementedError when Dataset.map etc is called

	try:
		yield
	finally:
		lowerCamelCase__              =	None



