


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


__A	:	int     =     datasets.utils.logging.get_logger(__name__)
@dataclass
class 						__UpperCamelCase	(						datasets.BuilderConfig							):
							lowercase			:		Optional[datasets.Features]   		=  None
							lowercase			:		str   		=  "utf-8"
							lowercase			:		Optional[str]   		=  None
							lowercase			:		Optional[str]   		=  None
							lowercase			:		bool   		=  True  # deprecated
							lowercase			:		Optional[int]   		=  None  # deprecated
							lowercase			:		int   		=  1_0 << 2_0  # 10MB
							lowercase			:		Optional[bool]   		=  None





class 						__UpperCamelCase	(						datasets.ArrowBasedBuilder							):
							lowercase			:		Tuple   		=  JsonConfig





							def 		a__   (  self		:Optional[int]			):
									if self.config.block_size is not None:
											logger.warning("""The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead"""			)
											snake_case_		:							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		:Optional[Any]  ,_UpperCamelCase		:Optional[int]			):
									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}'''			)
									snake_case_		:							Union[str, Any]			       =						dl_manager.download_and_extract(self.config.data_files			)
									if isinstance(_UpperCamelCase  ,(str, list, tuple)			):
											snake_case_		:							List[Any]			       =						data_files
											if isinstance(_UpperCamelCase  ,_UpperCamelCase			):
													snake_case_		:							int			       =						[files]
											snake_case_		:							Optional[Any]			       =						[dl_manager.iter_files(_UpperCamelCase			) for file in files]
											return [datasets.SplitGenerator(name=datasets.Split.TRAIN  ,gen_kwargs={"""files""": files}			)]
									snake_case_		:							int			       =						[]
									for split_name, files in data_files.items():
											if isinstance(_UpperCamelCase  ,_UpperCamelCase			):
													snake_case_		:							List[Any]			       =						[files]
											snake_case_		:							Optional[int]			       =						[dl_manager.iter_files(_UpperCamelCase			) for file in files]
											splits.append(datasets.SplitGenerator(name=_UpperCamelCase  ,gen_kwargs={"""files""": files}			)			)
									return splits





							def 		a__   (  self		:Union[str, Any]  ,_UpperCamelCase		: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			):
													snake_case_		:							Tuple			       =						self.config.features.arrow_schema.field(_UpperCamelCase			).type
													snake_case_		:							Optional[Any]			       =						pa_table.append_column(_UpperCamelCase  ,pa.array([None] * len(_UpperCamelCase			)  ,type=_UpperCamelCase			)			)
											# more expensive cast to support nested structures with keys in a different order
											# allows str <-> int/float or str to Audio for example
											snake_case_		:							Optional[Any]			       =						table_cast(_UpperCamelCase  ,self.config.features.arrow_schema			)
									return pa_table





							def 		a__   (  self		:List[Any]  ,_UpperCamelCase		:int			):
									for file_idx, file in enumerate(itertools.chain.from_iterable(_UpperCamelCase			)			):
											# 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(_UpperCamelCase  ,encoding=self.config.encoding  ,errors=self.config.encoding_errors			) as f:
															snake_case_		:							Union[str, Any]			       =						json.load(_UpperCamelCase			)

													# We keep only the field we are interested in
													snake_case_		:							Tuple			       =						dataset[self.config.field]

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

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



'''simple docstring'''






import os
import unittest

from transformers.models.transfo_xl.tokenization_transfo_xl import VOCAB_FILES_NAMES, TransfoXLTokenizer

from ...test_tokenization_common import TokenizerTesterMixin
class 						__UpperCamelCase	(						lowercase__						,  unittest.TestCase							):
							lowercase			:		Dict   		=  TransfoXLTokenizer
							lowercase			:		Optional[Any]   		=  False
							lowercase			:		Dict   		=  False





							def 		a__   (  self		:Union[str, Any]			):
									super().setUp()

									snake_case_		:							Optional[int]			       =						[
									    """<unk>""",
									    """[CLS]""",
									    """[SEP]""",
									    """want""",
									    """unwanted""",
									    """wa""",
									    """un""",
									    """running""",
									    """,""",
									    """low""",
									    """l""",
									]
									snake_case_		:							Optional[int]			       =						os.path.join(self.tmpdirname  ,VOCAB_FILES_NAMES["""vocab_file"""]			)
									with open(self.vocab_file  ,"""w"""  ,encoding="""utf-8"""			) as vocab_writer:
											vocab_writer.write("""""".join([x + """\n""" for x in vocab_tokens]			)			)





							def 		a__   (  self		:List[Any]  ,**_UpperCamelCase		:Optional[Any]			):
									snake_case_		:							Tuple			       =						True
									return TransfoXLTokenizer.from_pretrained(self.tmpdirname  ,**_UpperCamelCase			)





							def 		a__   (  self		:Tuple  ,_UpperCamelCase		:Union[str, Any]			):
									snake_case_		:							Any			       =						"""<unk> UNwanted , running"""
									snake_case_		:							Optional[int]			       =						"""<unk> unwanted, running"""
									return input_text, output_text





							def 		a__   (  self		:Dict			):
									snake_case_		:							Dict			       =						TransfoXLTokenizer(vocab_file=self.vocab_file  ,lower_case=_UpperCamelCase			)

									snake_case_		:							Dict			       =						tokenizer.tokenize("""<unk> UNwanted , running"""			)
									self.assertListEqual(_UpperCamelCase  ,["""<unk>""", """unwanted""", """,""", """running"""]			)

									self.assertListEqual(tokenizer.convert_tokens_to_ids(_UpperCamelCase			)  ,[0, 4, 8, 7]			)





							def 		a__   (  self		:Optional[Any]			):
									snake_case_		:							Dict			       =						TransfoXLTokenizer(lower_case=_UpperCamelCase			)

									self.assertListEqual(
									    tokenizer.tokenize(""" \tHeLLo ! how  \n Are yoU ?  """			)  ,["""hello""", """!""", """how""", """are""", """you""", """?"""]			)





							def 		a__   (  self		:Any			):
									snake_case_		:							List[Any]			       =						TransfoXLTokenizer(lower_case=_UpperCamelCase			)

									self.assertListEqual(
									    tokenizer.tokenize(""" \tHeLLo ! how  \n Are yoU ?  """			)  ,["""HeLLo""", """!""", """how""", """Are""", """yoU""", """?"""]			)





							def 		a__   (  self		:List[str]			):
									snake_case_		:							str			       =						TransfoXLTokenizer(lower_case=_UpperCamelCase			)
									snake_case_		:							List[str]			       =						"""Hello (bracket) and side-scrolled [and] Henry's $5,000 with 3.34 m. What's up!?"""
									snake_case_		:							Optional[int]			       =						[
									    """Hello""",
									    """(""",
									    """bracket""",
									    """)""",
									    """and""",
									    """side""",
									    """@-@""",
									    """scrolled""",
									    """[""",
									    """and""",
									    """]""",
									    """Henry""",
									    """'s""",
									    """$""",
									    """5""",
									    """@,@""",
									    """000""",
									    """with""",
									    """3""",
									    """@.@""",
									    """34""",
									    """m""",
									    """.""",
									    """What""",
									    """'s""",
									    """up""",
									    """!""",
									    """?""",
									]

									self.assertListEqual(tokenizer.tokenize(_UpperCamelCase			)  ,_UpperCamelCase			)

									self.assertEqual(tokenizer.convert_tokens_to_string(_UpperCamelCase			)  ,_UpperCamelCase			)





							def 		a__   (  self		:Dict			):
									snake_case_		:							Union[str, Any]			       =						self.get_tokenizer()
									snake_case_		:							Dict			       =						len(_UpperCamelCase			)

									tokenizer.add_tokens(["""new1""", """new2"""]			)
									tokenizer.move_added_token("""new1"""  ,1			)

									# Check that moved token is not copied (duplicate)
									self.assertEqual(len(_UpperCamelCase			)  ,original_len + 2			)
									# Check that token is moved to specified id
									self.assertEqual(tokenizer.encode("""new1"""			)  ,[1]			)
									self.assertEqual(tokenizer.decode([1]			)  ,"""new1"""			)