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

from tokenizers import processors

from ...tokenization_utils import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
							from .tokenization_nllb import NllbTokenizer
else:
							UpperCamelCase_						:     str	       =							None


UpperCamelCase_						:     List[str]	       =							logging.get_logger(__name__)


UpperCamelCase_						:     Dict	       =							{"""vocab_file""": """sentencepiece.bpe.model""", """tokenizer_file""": """tokenizer.json"""}

UpperCamelCase_						:     Union[str, Any]	       =							{
    """vocab_file""": {
        """facebook/nllb-200-distilled-600M""": (
            """https://huggingface.co/facebook/nllb-200-distilled-600M/resolve/main/sentencepiece.bpe.model"""
        ),
    },
    """tokenizer_file""": {
        """facebook/nllb-200-distilled-600M""": (
            """https://huggingface.co/facebook/nllb-200-distilled-600M/resolve/main/tokenizer.json"""
        ),
    },
}

UpperCamelCase_						:     Optional[Any]	       =							{
    """facebook/nllb-large-en-ro""": 1024,
    """facebook/nllb-200-distilled-600M""": 1024,
}

# fmt: off
UpperCamelCase_						:     List[Any]	       =							["""ace_Arab""", """ace_Latn""", """acm_Arab""", """acq_Arab""", """aeb_Arab""", """afr_Latn""", """ajp_Arab""", """aka_Latn""", """amh_Ethi""", """apc_Arab""", """arb_Arab""", """ars_Arab""", """ary_Arab""", """arz_Arab""", """asm_Beng""", """ast_Latn""", """awa_Deva""", """ayr_Latn""", """azb_Arab""", """azj_Latn""", """bak_Cyrl""", """bam_Latn""", """ban_Latn""", """bel_Cyrl""", """bem_Latn""", """ben_Beng""", """bho_Deva""", """bjn_Arab""", """bjn_Latn""", """bod_Tibt""", """bos_Latn""", """bug_Latn""", """bul_Cyrl""", """cat_Latn""", """ceb_Latn""", """ces_Latn""", """cjk_Latn""", """ckb_Arab""", """crh_Latn""", """cym_Latn""", """dan_Latn""", """deu_Latn""", """dik_Latn""", """dyu_Latn""", """dzo_Tibt""", """ell_Grek""", """eng_Latn""", """epo_Latn""", """est_Latn""", """eus_Latn""", """ewe_Latn""", """fao_Latn""", """pes_Arab""", """fij_Latn""", """fin_Latn""", """fon_Latn""", """fra_Latn""", """fur_Latn""", """fuv_Latn""", """gla_Latn""", """gle_Latn""", """glg_Latn""", """grn_Latn""", """guj_Gujr""", """hat_Latn""", """hau_Latn""", """heb_Hebr""", """hin_Deva""", """hne_Deva""", """hrv_Latn""", """hun_Latn""", """hye_Armn""", """ibo_Latn""", """ilo_Latn""", """ind_Latn""", """isl_Latn""", """ita_Latn""", """jav_Latn""", """jpn_Jpan""", """kab_Latn""", """kac_Latn""", """kam_Latn""", """kan_Knda""", """kas_Arab""", """kas_Deva""", """kat_Geor""", """knc_Arab""", """knc_Latn""", """kaz_Cyrl""", """kbp_Latn""", """kea_Latn""", """khm_Khmr""", """kik_Latn""", """kin_Latn""", """kir_Cyrl""", """kmb_Latn""", """kon_Latn""", """kor_Hang""", """kmr_Latn""", """lao_Laoo""", """lvs_Latn""", """lij_Latn""", """lim_Latn""", """lin_Latn""", """lit_Latn""", """lmo_Latn""", """ltg_Latn""", """ltz_Latn""", """lua_Latn""", """lug_Latn""", """luo_Latn""", """lus_Latn""", """mag_Deva""", """mai_Deva""", """mal_Mlym""", """mar_Deva""", """min_Latn""", """mkd_Cyrl""", """plt_Latn""", """mlt_Latn""", """mni_Beng""", """khk_Cyrl""", """mos_Latn""", """mri_Latn""", """zsm_Latn""", """mya_Mymr""", """nld_Latn""", """nno_Latn""", """nob_Latn""", """npi_Deva""", """nso_Latn""", """nus_Latn""", """nya_Latn""", """oci_Latn""", """gaz_Latn""", """ory_Orya""", """pag_Latn""", """pan_Guru""", """pap_Latn""", """pol_Latn""", """por_Latn""", """prs_Arab""", """pbt_Arab""", """quy_Latn""", """ron_Latn""", """run_Latn""", """rus_Cyrl""", """sag_Latn""", """san_Deva""", """sat_Beng""", """scn_Latn""", """shn_Mymr""", """sin_Sinh""", """slk_Latn""", """slv_Latn""", """smo_Latn""", """sna_Latn""", """snd_Arab""", """som_Latn""", """sot_Latn""", """spa_Latn""", """als_Latn""", """srd_Latn""", """srp_Cyrl""", """ssw_Latn""", """sun_Latn""", """swe_Latn""", """swh_Latn""", """szl_Latn""", """tam_Taml""", """tat_Cyrl""", """tel_Telu""", """tgk_Cyrl""", """tgl_Latn""", """tha_Thai""", """tir_Ethi""", """taq_Latn""", """taq_Tfng""", """tpi_Latn""", """tsn_Latn""", """tso_Latn""", """tuk_Latn""", """tum_Latn""", """tur_Latn""", """twi_Latn""", """tzm_Tfng""", """uig_Arab""", """ukr_Cyrl""", """umb_Latn""", """urd_Arab""", """uzn_Latn""", """vec_Latn""", """vie_Latn""", """war_Latn""", """wol_Latn""", """xho_Latn""", """ydd_Hebr""", """yor_Latn""", """yue_Hant""", """zho_Hans""", """zho_Hant""", """zul_Latn"""]




class   __lowercase						(						__snake_case       ):
	_A       =   VOCAB_FILES_NAMES
	_A       =   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	_A       =   PRETRAINED_VOCAB_FILES_MAP
	_A       =   ["input_ids", "attention_mask"]
	_A       =   NllbTokenizer

	_A       =   []
	_A       =   []






	def __init__(self				: List[Any]      ,	snake_case				: Tuple=None      ,	snake_case				: Union[str, Any]=None      ,	snake_case				: Optional[int]="<s>"      ,	snake_case				: str="</s>"      ,	snake_case				: Optional[int]="</s>"      ,	snake_case				: Tuple="<s>"      ,	snake_case				: int="<unk>"      ,	snake_case				: Dict="<pad>"      ,	snake_case				: List[Any]="<mask>"      ,	snake_case				: Dict=None      ,	snake_case				: Any=None      ,	snake_case				: str=None      ,	snake_case				: Union[str, Any]=False      ,	**snake_case				: Any      ,	) ->			Tuple:
							# Mask token behave like a normal word, i.e. include the space before it
							_lowercase       :	int           =   AddedToken(snake_case      ,	lstrip=snake_case      ,	rstrip=snake_case    ) if isinstance(snake_case      ,	snake_case    ) else mask_token
							_lowercase       :	Optional[Any]           =   legacy_behaviour
							super().__init__(
							    vocab_file=snake_case      ,	tokenizer_file=snake_case      ,	bos_token=snake_case      ,	eos_token=snake_case      ,	sep_token=snake_case      ,	cls_token=snake_case      ,	unk_token=snake_case      ,	pad_token=snake_case      ,	mask_token=snake_case      ,	src_lang=snake_case      ,	tgt_lang=snake_case      ,	additional_special_tokens=snake_case      ,	legacy_behaviour=snake_case      ,	**snake_case      ,	)

							_lowercase       :	int           =   vocab_file
							_lowercase       :	Optional[Any]           =   False if not self.vocab_file else True

							_lowercase       :	List[Any]           =   FAIRSEQ_LANGUAGE_CODES.copy()

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

							self.add_special_tokens({"additional_special_tokens": _additional_special_tokens}    )
							_lowercase       :	Union[str, Any]           =   {
							    lang_code: self.convert_tokens_to_ids(snake_case    ) for lang_code in FAIRSEQ_LANGUAGE_CODES
							}

							_lowercase       :	str           =   src_lang if src_lang is not None else "eng_Latn"
							_lowercase       :	Any           =   self.convert_tokens_to_ids(self._src_lang    )
							_lowercase       :	Union[str, Any]           =   tgt_lang
							self.set_src_lang_special_tokens(self._src_lang    )






	@property
	def       _a(self				: Any    ) ->			str:
							return self._src_lang






	@src_lang.setter
	def       _a(self				: Dict      ,	snake_case				: str    ) ->			None:
							_lowercase       :	Optional[int]           =   new_src_lang
							self.set_src_lang_special_tokens(self._src_lang    )






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






	def       _a(self				: int      ,	snake_case				: List[str]      ,	snake_case				: str = "eng_Latn"      ,	snake_case				: Optional[List[str]] = None      ,	snake_case				: str = "fra_Latn"      ,	**snake_case				: Optional[Any]      ,	) ->			BatchEncoding:
							_lowercase       :	Optional[int]           =   src_lang
							_lowercase       :	Any           =   tgt_lang
							return super().prepare_seqaseq_batch(snake_case      ,	snake_case      ,	**snake_case    )






	def       _a(self				: List[str]    ) ->			List[str]:
							return self.set_src_lang_special_tokens(self.src_lang    )






	def       _a(self				: List[Any]    ) ->			List[Any]:
							return self.set_tgt_lang_special_tokens(self.tgt_lang    )






	def       _a(self				: List[str]      ,	snake_case				: List[str]    ) ->			None:
							_lowercase       :	List[Any]           =   self.convert_tokens_to_ids(snake_case    )

							if self.legacy_behaviour:
													_lowercase       :	Union[str, Any]           =   []
													_lowercase       :	Optional[Any]           =   [self.eos_token_id, self.cur_lang_code]
							else:
													_lowercase       :	Dict           =   [self.cur_lang_code]
													_lowercase       :	Tuple           =   [self.eos_token_id]

							_lowercase       :	List[Any]           =   self.convert_ids_to_tokens(self.prefix_tokens    )
							_lowercase       :	Any           =   self.convert_ids_to_tokens(self.suffix_tokens    )

							_lowercase       :	Union[str, Any]           =   processors.TemplateProcessing(
							    single=prefix_tokens_str + ["$A"] + suffix_tokens_str      ,	pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str      ,	special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str      ,	self.prefix_tokens + self.suffix_tokens    )    )      ,	)






	def       _a(self				: Dict      ,	snake_case				: str    ) ->			None:
							_lowercase       :	int           =   self.convert_tokens_to_ids(snake_case    )
							if self.legacy_behaviour:
													_lowercase       :	Optional[Any]           =   []
													_lowercase       :	int           =   [self.eos_token_id, self.cur_lang_code]
							else:
													_lowercase       :	Any           =   [self.cur_lang_code]
													_lowercase       :	int           =   [self.eos_token_id]

							_lowercase       :	Optional[Any]           =   self.convert_ids_to_tokens(self.prefix_tokens    )
							_lowercase       :	List[str]           =   self.convert_ids_to_tokens(self.suffix_tokens    )

							_lowercase       :	List[Any]           =   processors.TemplateProcessing(
							    single=prefix_tokens_str + ["$A"] + suffix_tokens_str      ,	pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str      ,	special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str      ,	self.prefix_tokens + self.suffix_tokens    )    )      ,	)






	def       _a(self				: Dict      ,	snake_case				: str      ,	snake_case				: Optional[str] = None    ) ->			Tuple[str]:
							if not self.can_save_slow_tokenizer:
													raise ValueError(
													    "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
													    "tokenizer."    )

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

							if os.path.abspath(self.vocab_file    ) != os.path.abspath(snake_case    ):
													copyfile(self.vocab_file      ,	snake_case    )

							return (out_vocab_file,)






import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


UpperCamelCase_						:     Union[str, Any]	       =							datasets.utils.logging.get_logger(__name__)




@dataclass
class   __lowercase						(						datasets.BuilderConfig       ):
	_A       =   None
	_A       =   "utf-8"
	_A       =   None
	_A       =   None
	_A       =   True  # deprecated
	_A       =   None  # deprecated
	_A       =   10 << 20  # 10MB
	_A       =   None




class   __lowercase						(						datasets.ArrowBasedBuilder       ):
	_A       =   JsonConfig






	def       _a(self				: List[Any]    ) ->			Dict:
							if self.config.block_size is not None:
													logger.warning("The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead"    )
													_lowercase       :	Union[str, Any]           =   self.config.block_size
							if self.config.use_threads is not True:
													logger.warning(
													    "The JSON loader parameter `use_threads` is deprecated and doesn't have any effect anymore."    )
							if self.config.newlines_in_values is not None:
													raise ValueError("The JSON loader parameter `newlines_in_values` is no longer supported"    )
							return datasets.DatasetInfo(features=self.config.features    )






	def       _a(self				: Any      ,	snake_case				: Union[str, Any]    ) ->			List[str]:
							if not self.config.data_files:
													raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}"""    )
							_lowercase       :	str           =   dl_manager.download_and_extract(self.config.data_files    )
							if isinstance(snake_case      ,	(str, list, tuple)    ):
													_lowercase       :	int           =   data_files
													if isinstance(snake_case      ,	snake_case    ):
																			_lowercase       :	Union[str, Any]           =   [files]
													_lowercase       :	Union[str, Any]           =   [dl_manager.iter_files(snake_case    ) for file in files]
													return [datasets.SplitGenerator(name=datasets.Split.TRAIN      ,	gen_kwargs={"files": files}    )]
							_lowercase       :	Union[str, Any]           =   []
							for split_name, files in data_files.items():
													if isinstance(snake_case      ,	snake_case    ):
																			_lowercase       :	List[Any]           =   [files]
													_lowercase       :	Optional[Any]           =   [dl_manager.iter_files(snake_case    ) for file in files]
													splits.append(datasets.SplitGenerator(name=snake_case      ,	gen_kwargs={"files": files}    )    )
							return splits






	def       _a(self				: List[Any]      ,	snake_case				: pa.Table    ) ->			pa.Table:
							if self.config.features is not None:
													# adding missing columns
													for column_name in set(self.config.features    ) - set(pa_table.column_names    ):
																			_lowercase       :	Any           =   self.config.features.arrow_schema.field(snake_case    ).type
																			_lowercase       :	Optional[Any]           =   pa_table.append_column(snake_case      ,	pa.array([None] * len(snake_case    )      ,	type=snake_case    )    )
													# more expensive cast to support nested structures with keys in a different order
													# allows str <-> int/float or str to Audio for example
													_lowercase       :	str           =   table_cast(snake_case      ,	self.config.features.arrow_schema    )
							return pa_table






	def       _a(self				: str      ,	snake_case				: Tuple    ) ->			Tuple:
							for file_idx, file in enumerate(itertools.chain.from_iterable(snake_case    )    ):
													# If the file is one json object and if we need to look at the list of items in one specific field
													if self.config.field is not None:
																			with open(snake_case      ,	encoding=self.config.encoding      ,	errors=self.config.encoding_errors    ) as f:
																									_lowercase       :	Any           =   json.load(snake_case    )

																			# We keep only the field we are interested in
																			_lowercase       :	List[Any]           =   dataset[self.config.field]

																			# We accept two format: a list of dicts or a dict of lists
																			if isinstance(snake_case      ,	(list, tuple)    ):
																									_lowercase       :	Optional[int]           =   set().union(*[row.keys() for row in dataset]    )
																									_lowercase       :	Union[str, Any]           =   {col: [row.get(snake_case    ) for row in dataset] for col in keys}
																			else:
																									_lowercase       :	List[Any]           =   dataset
																			_lowercase       :	int           =   pa.Table.from_pydict(snake_case    )
																			yield file_idx, self._cast_table(snake_case    )

													# If the file has one json object per line
													else:
																			with open(snake_case      ,	"rb"    ) as f:
																									_lowercase       :	Union[str, Any]           =   0
																									# Use block_size equal to the chunk size divided by 32 to leverage multithreading
																									# Set a default minimum value of 16kB if the chunk size is really small
																									_lowercase       :	Dict           =   max(self.config.chunksize // 32      ,	16 << 10    )
																									_lowercase       :	int           =   (
																									    self.config.encoding_errors if self.config.encoding_errors is not None else "strict"
																									)
																									while True:
																															_lowercase       :	Tuple           =   f.read(self.config.chunksize    )
																															if not batch:
																																					break
																															# Finish current line
																															try:
																																					batch += f.readline()
																															except (AttributeError, io.UnsupportedOperation):
																																					batch += readline(snake_case    )
																															# PyArrow only accepts utf-8 encoded bytes
																															if self.config.encoding != "utf-8":
																																					_lowercase       :	Union[str, Any]           =   batch.decode(self.config.encoding      ,	errors=snake_case    ).encode("utf-8"    )
																															try:
																																					while True:
																																											try:
																																																	_lowercase       :	Optional[Any]           =   paj.read_json(
																																																	    io.BytesIO(snake_case    )      ,	read_options=paj.ReadOptions(block_size=snake_case    )    )
																																																	break
																																											except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
																																																	if (
																																																	    isinstance(snake_case      ,	pa.ArrowInvalid    )
																																																	    and "straddling" not in str(snake_case    )
																																																	    or block_size > len(snake_case    )
																																																	):
																																																							raise
																																																	else:
																																																							# Increase the block size in case it was too small.
																																																							# The block size will be reset for the next file.
																																																							logger.debug(
																																																							    F"""Batch of {len(snake_case    )} bytes couldn't be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}."""    )
																																																							block_size *= 2
																															except pa.ArrowInvalid as e:
																																					try:
																																											with open(
																																											    snake_case      ,	encoding=self.config.encoding      ,	errors=self.config.encoding_errors    ) as f:
																																																	_lowercase       :	Any           =   json.load(snake_case    )
																																					except json.JSONDecodeError:
																																											logger.error(F"""Failed to read file '{file}' with error {type(snake_case    )}: {e}"""    )
																																											raise e
																																					# If possible, parse the file as a list of json objects and exit the loop
																																					if isinstance(snake_case      ,	snake_case    ):  # list is the only sequence type supported in JSON
																																											try:
																																																	_lowercase       :	int           =   set().union(*[row.keys() for row in dataset]    )
																																																	_lowercase       :	Dict           =   {col: [row.get(snake_case    ) for row in dataset] for col in keys}
																																																	_lowercase       :	str           =   pa.Table.from_pydict(snake_case    )
																																											except (pa.ArrowInvalid, AttributeError) as e:
																																																	logger.error(F"""Failed to read file '{file}' with error {type(snake_case    )}: {e}"""    )
																																																	raise ValueError(F"""Not able to read records in the JSON file at {file}."""    ) from None
																																											yield file_idx, self._cast_table(snake_case    )
																																											break
																																					else:
																																											logger.error(F"""Failed to read file '{file}' with error {type(snake_case    )}: {e}"""    )
																																											raise ValueError(
																																											    F"""Not able to read records in the JSON file at {file}. """
																																											    F"""You should probably indicate the field of the JSON file containing your records. """
																																											    F"""This JSON file contain the following fields: {str(list(dataset.keys()    )    )}. """
																																											    F"""Select the correct one and provide it as `field='XXX'` to the dataset loading method. """    ) from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																															yield (file_idx, batch_idx), self._cast_table(snake_case    )
																															batch_idx += 1




