

'''simple docstring'''





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_							=					datasets.utils.logging.get_logger(__name__)





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





class 		__SCREAMING_SNAKE_CASE       ( datasets.ArrowBasedBuilder			):
				lowerCamelCase_							=						JsonConfig







				def    lowerCamelCase_  (					self    :		Union[str, Any]      ):
										'''simple docstring'''
										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    lowerCamelCase_  (					self    :		Tuple		,				UpperCAmelCase__    :		Any      ):
										'''simple docstring'''
										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   :       int								=dl_manager.download_and_extract(self.config.data_files      )
										if isinstance(A__		,				(str, list, tuple)      ):
																lowercase   :       Optional[Any]								=data_files
																if isinstance(A__		,				A__      ):
																						lowercase   :       List[str]								=[files]
																lowercase   :       int								=[dl_manager.iter_files(A__      ) for file in files]
																return [datasets.SplitGenerator(name=datasets.Split.TRAIN		,				gen_kwargs={'''files''': files}      )]
										lowercase   :       List[str]								=[]
										for split_name, files in data_files.items():
																if isinstance(A__		,				A__      ):
																						lowercase   :       Optional[int]								=[files]
																lowercase   :       Optional[int]								=[dl_manager.iter_files(A__      ) for file in files]
																splits.append(datasets.SplitGenerator(name=A__		,				gen_kwargs={'''files''': files}      )      )
										return splits







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







				def    lowerCamelCase_  (					self    :		Optional[int]		,				UpperCAmelCase__    :		Union[str, Any]      ):
										'''simple docstring'''
										for file_idx, file in enumerate(itertools.chain.from_iterable(A__      )      ):
																# 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(A__		,				encoding=self.config.encoding		,				errors=self.config.encoding_errors      ) as f:
																												lowercase   :       Optional[Any]								=json.load(A__      )

																						# We keep only the field we are interested in
																						lowercase   :       Optional[int]								=dataset[self.config.field]

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

																# If the file has one json object per line
																else:
																						with open(A__		,				'''rb'''      ) as f:
																												lowercase   :       List[str]								=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   :       List[str]								=max(self.config.chunksize // 32		,				16 << 10      )
																												lowercase   :       Any								=(
																												    self.config.encoding_errors if self.config.encoding_errors is not None else """strict"""
																												)
																												while True:
																																		lowercase   :       Dict								=f.read(self.config.chunksize      )
																																		if not batch:
																																								break
																																		# Finish current line
																																		try:
																																								batch += f.readline()
																																		except (AttributeError, io.UnsupportedOperation):
																																								batch += readline(A__      )
																																		# PyArrow only accepts utf-8 encoded bytes
																																		if self.config.encoding != "utf-8":
																																								lowercase   :       List[Any]								=batch.decode(self.config.encoding		,				errors=A__      ).encode('''utf-8'''      )
																																		try:
																																								while True:
																																														try:
																																																				lowercase   :       str								=paj.read_json(
																																																				    io.BytesIO(A__      )		,				read_options=paj.ReadOptions(block_size=A__      )      )
																																																				break
																																														except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
																																																				if (
																																																				    isinstance(A__		,				pa.ArrowInvalid      )
																																																				    and "straddling" not in str(A__      )
																																																				    or block_size > len(A__      )
																																																				):
																																																										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(A__      )} 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(
																																														    A__		,				encoding=self.config.encoding		,				errors=self.config.encoding_errors      ) as f:
																																																				lowercase   :       Optional[Any]								=json.load(A__      )
																																								except json.JSONDecodeError:
																																														logger.error(F'''Failed to read file \'{file}\' with error {type(A__      )}: {e}'''      )
																																														raise e
																																								# If possible, parse the file as a list of json objects and exit the loop
																																								if isinstance(A__		,				A__      ):  # list is the only sequence type supported in JSON
																																														try:
																																																				lowercase   :       str								=set().union(*[row.keys() for row in dataset]      )
																																																				lowercase   :       List[str]								={col: [row.get(A__      ) for row in dataset] for col in keys}
																																																				lowercase   :       int								=pa.Table.from_pydict(A__      )
																																														except (pa.ArrowInvalid, AttributeError) as e:
																																																				logger.error(F'''Failed to read file \'{file}\' with error {type(A__      )}: {e}'''      )
																																																				raise ValueError(F'''Not able to read records in the JSON file at {file}.'''      ) from None
																																														yield file_idx, self._cast_table(A__      )
																																														break
																																								else:
																																														logger.error(F'''Failed to read file \'{file}\' with error {type(A__      )}: {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(A__      )
																																		batch_idx += 1







'''simple docstring'''





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

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


UpperCamelCase_							=					logging.get_logger(__name__)

UpperCamelCase_							=					"""▁"""

UpperCamelCase_							=					{"""vocab_file""": """sentencepiece.bpe.model"""}

UpperCamelCase_							=					{
    """vocab_file""": {
        """facebook/xglm-564M""": """https://huggingface.co/facebook/xglm-564M/resolve/main/sentencepiece.bpe.model""",
    }
}

UpperCamelCase_							=					{
    """facebook/xglm-564M""": 2048,
}





class 		__SCREAMING_SNAKE_CASE       ( lowercase__			):
				lowerCamelCase_							=						VOCAB_FILES_NAMES
				lowerCamelCase_							=						PRETRAINED_VOCAB_FILES_MAP
				lowerCamelCase_							=						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				lowerCamelCase_							=						['input_ids', 'attention_mask']







				def __init__(					self    :		Dict		,				UpperCAmelCase__    :		Optional[int]		,				UpperCAmelCase__    :		List[Any]="<s>"		,				UpperCAmelCase__    :		int="</s>"		,				UpperCAmelCase__    :		Optional[Any]="</s>"		,				UpperCAmelCase__    :		Optional[Any]="<s>"		,				UpperCAmelCase__    :		Any="<unk>"		,				UpperCAmelCase__    :		Any="<pad>"		,				UpperCAmelCase__    :		Optional[Dict[str, Any]] = None		,				**UpperCAmelCase__    :		Optional[int]		,				):
										'''simple docstring'''
										lowercase   :       int								={} if sp_model_kwargs is None else sp_model_kwargs

										# Compatibility with the original tokenizer
										lowercase   :       Optional[Any]								=7
										lowercase   :       Optional[int]								=[F'''<madeupword{i}>''' for i in range(self.num_madeup_words      )]

										lowercase   :       List[Any]								=kwargs.get('''additional_special_tokens'''		,				[]      )
										kwargs["additional_special_tokens"] += [
										    word for word in madeup_words if word not in kwargs["additional_special_tokens"]
										]

										super().__init__(
										    bos_token=UpperCAmelCase__		,				eos_token=UpperCAmelCase__		,				unk_token=UpperCAmelCase__		,				sep_token=UpperCAmelCase__		,				cls_token=UpperCAmelCase__		,				pad_token=UpperCAmelCase__		,				sp_model_kwargs=self.sp_model_kwargs		,				**UpperCAmelCase__		,				)

										lowercase   :       List[str]								=spm.SentencePieceProcessor(**self.sp_model_kwargs      )
										self.sp_model.Load(str(UpperCAmelCase__      )      )
										lowercase   :       List[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'

										# The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
										lowercase   :       Union[str, Any]								=1

										# Mimic fairseq token-to-id alignment for the first 4 token
										lowercase   :       List[str]								={'''<s>''': 0, '''<pad>''': 1, '''</s>''': 2, '''<unk>''': 3}

										lowercase   :       str								=len(self.sp_model      )
										lowercase   :       List[Any]								={F'''<madeupword{i}>''': sp_size + i + self.fairseq_offset for i in range(self.num_madeup_words      )}
										self.fairseq_tokens_to_ids.update(UpperCAmelCase__      )

										lowercase   :       int								={v: k for k, v in self.fairseq_tokens_to_ids.items()}







				def __getstate__(					self    :		int      ):
										'''simple docstring'''
										lowercase   :       Optional[int]								=self.__dict__.copy()
										lowercase   :       List[Any]								=None
										lowercase   :       Tuple								=self.sp_model.serialized_model_proto()
										return state







				def __setstate__(					self    :		Optional[Any]		,				UpperCAmelCase__    :		Any      ):
										'''simple docstring'''
										lowercase   :       int								=d

										# for backward compatibility
										if not hasattr(self		,				'''sp_model_kwargs'''      ):
																lowercase   :       Optional[int]								={}

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







				def    lowerCamelCase_  (					self    :		str		,				UpperCAmelCase__    :		List[int]		,				UpperCAmelCase__    :		Optional[List[int]] = None      ):
										'''simple docstring'''

										if token_ids_a is None:
																return [self.sep_token_id] + token_ids_a
										lowercase   :       List[Any]								=[self.sep_token_id]
										return sep + token_ids_a + sep + sep + token_ids_a







				def    lowerCamelCase_  (					self    :		Any		,				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__      )

										if token_ids_a is None:
																return [1] + ([0] * len(UpperCAmelCase__      ))
										return [1] + ([0] * len(UpperCAmelCase__      )) + [1, 1] + ([0] * len(UpperCAmelCase__      ))







				def    lowerCamelCase_  (					self    :		Dict		,				UpperCAmelCase__    :		List[int]		,				UpperCAmelCase__    :		Optional[List[int]] = None      ):
										'''simple docstring'''
										lowercase   :       int								=[self.sep_token_id]

										if token_ids_a is None:
																return len(sep + token_ids_a      ) * [0]
										return len(sep + token_ids_a + sep + sep + token_ids_a      ) * [0]







				@property
				def    lowerCamelCase_  (					self    :		Dict      ):
										'''simple docstring'''
										return len(self.sp_model      ) + self.fairseq_offset + self.num_madeup_words







				def    lowerCamelCase_  (					self    :		List[str]      ):
										'''simple docstring'''
										lowercase   :       int								={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    :		Tuple		,				UpperCAmelCase__    :		str      ):
										'''simple docstring'''
										return self.sp_model.encode(UpperCAmelCase__		,				out_type=UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		Dict		,				UpperCAmelCase__    :		Optional[Any]      ):
										'''simple docstring'''
										if token in self.fairseq_tokens_to_ids:
																return self.fairseq_tokens_to_ids[token]
										lowercase   :       List[str]								=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    :		Optional[Any]		,				UpperCAmelCase__    :		Any      ):
										'''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    :		List[str]		,				UpperCAmelCase__    :		Optional[Any]      ):
										'''simple docstring'''
										lowercase   :       Dict								=''''''.join(UpperCAmelCase__      ).replace(UpperCAmelCase__		,				''' '''      ).strip()
										return out_string







				def    lowerCamelCase_  (					self    :		Optional[int]		,				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
										lowercase   :       Dict								=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:
																						lowercase   :       Optional[int]								=self.sp_model.serialized_model_proto()
																						fi.write(UpperCAmelCase__      )

										return (out_vocab_file,)





