


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


lowerCamelCase_       :       Optional[int]      						=  datasets.utils.logging.get_logger(__name__)






@dataclass
class 				__A    (	datasets.BuilderConfig      ):






 """simple docstring"""




 __lowerCAmelCase			 =  None
 __lowerCAmelCase			 =  "utf-8"
 __lowerCAmelCase			 =  None
 __lowerCAmelCase			 =  None
 __lowerCAmelCase			 =  True  # deprecated
 __lowerCAmelCase			 =  None  # deprecated
 __lowerCAmelCase			 =  10 << 20  # 10MB
 __lowerCAmelCase			 =  None






class 				__A    (	datasets.ArrowBasedBuilder      ):






 """simple docstring"""




 __lowerCAmelCase			 =  JsonConfig


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


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A						) ->					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						):
             a					=self.config.features.arrow_schema.field(__A						).type
             a					=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
         a					=table_cast(__A			,   self.config.features.arrow_schema						)
     return pa_table





 def 				SCREAMING_SNAKE_CASE							(		self			,   __A						) ->					List[str]:
     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:
                 a					=json.load(__A						)

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

             # We accept two format: a list of dicts or a dict of lists
             if isinstance(__A			,   (list, tuple)						):
                 a					=set().union(*[row.keys() for row in dataset]						)
                 a					={col: [row.get(__A						) for row in dataset] for col in keys}
             else:
                 a					=dataset
             a					=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:
                 a					=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
                 a					=max(self.config.chunksize // 32			,   16 << 10						)
                 a					=(
                     self.config.encoding_errors if self.config.encoding_errors is not None else '''strict'''
                 )
                 while True:
                     a					=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":
                         a					=batch.decode(self.config.encoding			,   errors=__A						).encode('''utf-8'''						)
                     try:
                         while True:
                             try:
                                 a					=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:
                                 a					=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:
                                 a					=set().union(*[row.keys() for row in dataset]						)
                                 a					={col: [row.get(__A						) for row in dataset] for col in keys}
                                 a					=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"""







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowerCamelCase_       :       int      						=  {"""configuration_mbart""": ["""MBART_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MBartConfig""", """MBartOnnxConfig"""]}

try:
       if not is_sentencepiece_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase_       :       Dict      						=  ["""MBartTokenizer"""]

try:
       if not is_tokenizers_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase_       :       str      						=  ["""MBartTokenizerFast"""]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase_       :       List[str]      						=  [
           """MBART_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """MBartForCausalLM""",
           """MBartForConditionalGeneration""",
           """MBartForQuestionAnswering""",
           """MBartForSequenceClassification""",
           """MBartModel""",
           """MBartPreTrainedModel""",
       ]

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase_       :       List[Any]      						=  [
           """TFMBartForConditionalGeneration""",
           """TFMBartModel""",
           """TFMBartPreTrainedModel""",
       ]

try:
       if not is_flax_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCamelCase_       :       str      						=  [
           """FlaxMBartForConditionalGeneration""",
           """FlaxMBartForQuestionAnswering""",
           """FlaxMBartForSequenceClassification""",
           """FlaxMBartModel""",
           """FlaxMBartPreTrainedModel""",
       ]


if TYPE_CHECKING:
       from .configuration_mbart import MBART_PRETRAINED_CONFIG_ARCHIVE_MAP, MBartConfig, MBartOnnxConfig

       try:
              if not is_sentencepiece_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_mbart import MBartTokenizer

       try:
              if not is_tokenizers_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_mbart_fast import MBartTokenizerFast

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_mbart import (
                  MBART_PRETRAINED_MODEL_ARCHIVE_LIST,
                  MBartForCausalLM,
                  MBartForConditionalGeneration,
                  MBartForQuestionAnswering,
                  MBartForSequenceClassification,
                  MBartModel,
                  MBartPreTrainedModel,
              )

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_mbart import TFMBartForConditionalGeneration, TFMBartModel, TFMBartPreTrainedModel

       try:
              if not is_flax_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_flax_mbart import (
                  FlaxMBartForConditionalGeneration,
                  FlaxMBartForQuestionAnswering,
                  FlaxMBartForSequenceClassification,
                  FlaxMBartModel,
                  FlaxMBartPreTrainedModel,
              )

else:
       import sys

       lowerCamelCase_       :       List[Any]      						=  _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)