



"""simple docstring"""
import tempfile

import torch

from diffusers import PNDMScheduler

from .test_schedulers import SchedulerCommonTest



class       SCREAMING_SNAKE_CASE__  (					UpperCAmelCase__ ):
   __lowerCAmelCase  :   Dict						=						(PNDMScheduler,)
   __lowerCAmelCase  :   int						=						(('num_inference_steps', 50),)

   def 	SCREAMING_SNAKE_CASE     (   self      ,     **_SCREAMING_SNAKE_CASE						)  ->      Optional[Any]:





    '''simple docstring'''




    UpperCAmelCase	:					Tuple	=	{
        """num_train_timesteps""": 1000,
        """beta_start""": 0.0001,
        """beta_end""": 0.02,
        """beta_schedule""": """linear""",
    }

    config.update(**_SCREAMING_SNAKE_CASE						)
    return config

   def 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE=0      ,     **_SCREAMING_SNAKE_CASE						)  ->      int:





    '''simple docstring'''




    UpperCAmelCase	:					Optional[Any]	=	dict(self.forward_default_kwargs						)
    UpperCAmelCase	:					Any	=	kwargs.pop("""num_inference_steps"""      ,     _SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Optional[Any]	=	self.dummy_sample
    UpperCAmelCase	:					Dict	=	0.1 * sample
    UpperCAmelCase	:					Optional[int]	=	[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

    for scheduler_class in self.scheduler_classes:
     UpperCAmelCase	:					str	=	self.get_scheduler_config(**_SCREAMING_SNAKE_CASE						)
     UpperCAmelCase	:					List[Any]	=	scheduler_class(**_SCREAMING_SNAKE_CASE						)
     scheduler.set_timesteps(_SCREAMING_SNAKE_CASE						)
     # copy over dummy past residuals
     UpperCAmelCase	:					Tuple	=	dummy_past_residuals[:]

     with tempfile.TemporaryDirectory() as tmpdirname:
      scheduler.save_config(_SCREAMING_SNAKE_CASE						)
      UpperCAmelCase	:					Dict	=	scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE						)
      new_scheduler.set_timesteps(_SCREAMING_SNAKE_CASE						)
      # copy over dummy past residuals
      UpperCAmelCase	:					Union[str, Any]	=	dummy_past_residuals[:]

     UpperCAmelCase	:					int	=	scheduler.step_prk(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     **_SCREAMING_SNAKE_CASE						).prev_sample
     UpperCAmelCase	:					str	=	new_scheduler.step_prk(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     **_SCREAMING_SNAKE_CASE						).prev_sample

     assert torch.sum(torch.abs(output - new_output						)						) < 1E-5, "Scheduler outputs are not identical"

     UpperCAmelCase	:					int	=	scheduler.step_plms(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     **_SCREAMING_SNAKE_CASE						).prev_sample
     UpperCAmelCase	:					List[Any]	=	new_scheduler.step_plms(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     **_SCREAMING_SNAKE_CASE						).prev_sample

     assert torch.sum(torch.abs(output - new_output						)						) < 1E-5, "Scheduler outputs are not identical"

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      List[str]:





    '''simple docstring'''




    pass

   def 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE=0      ,     **_SCREAMING_SNAKE_CASE						)  ->      List[str]:





    '''simple docstring'''




    UpperCAmelCase	:					Union[str, Any]	=	dict(self.forward_default_kwargs						)
    UpperCAmelCase	:					Tuple	=	kwargs.pop("""num_inference_steps"""      ,     _SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Any	=	self.dummy_sample
    UpperCAmelCase	:					Tuple	=	0.1 * sample
    UpperCAmelCase	:					Union[str, Any]	=	[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

    for scheduler_class in self.scheduler_classes:
     UpperCAmelCase	:					str	=	self.get_scheduler_config()
     UpperCAmelCase	:					Tuple	=	scheduler_class(**_SCREAMING_SNAKE_CASE						)
     scheduler.set_timesteps(_SCREAMING_SNAKE_CASE						)

     # copy over dummy past residuals (must be after setting timesteps)
     UpperCAmelCase	:					Any	=	dummy_past_residuals[:]

     with tempfile.TemporaryDirectory() as tmpdirname:
      scheduler.save_config(_SCREAMING_SNAKE_CASE						)
      UpperCAmelCase	:					List[str]	=	scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE						)
      # copy over dummy past residuals
      new_scheduler.set_timesteps(_SCREAMING_SNAKE_CASE						)

      # copy over dummy past residual (must be after setting timesteps)
      UpperCAmelCase	:					Dict	=	dummy_past_residuals[:]

     UpperCAmelCase	:					Dict	=	scheduler.step_prk(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     **_SCREAMING_SNAKE_CASE						).prev_sample
     UpperCAmelCase	:					Optional[Any]	=	new_scheduler.step_prk(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     **_SCREAMING_SNAKE_CASE						).prev_sample

     assert torch.sum(torch.abs(output - new_output						)						) < 1E-5, "Scheduler outputs are not identical"

     UpperCAmelCase	:					Tuple	=	scheduler.step_plms(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     **_SCREAMING_SNAKE_CASE						).prev_sample
     UpperCAmelCase	:					List[Any]	=	new_scheduler.step_plms(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     **_SCREAMING_SNAKE_CASE						).prev_sample

     assert torch.sum(torch.abs(output - new_output						)						) < 1E-5, "Scheduler outputs are not identical"

   def 	SCREAMING_SNAKE_CASE     (   self      ,     **_SCREAMING_SNAKE_CASE						)  ->      List[Any]:





    '''simple docstring'''




    UpperCAmelCase	:					List[Any]	=	self.scheduler_classes[0]
    UpperCAmelCase	:					List[Any]	=	self.get_scheduler_config(**_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Tuple	=	scheduler_class(**_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Any	=	10
    UpperCAmelCase	:					Optional[int]	=	self.dummy_model()
    UpperCAmelCase	:					Any	=	self.dummy_sample_deter
    scheduler.set_timesteps(_SCREAMING_SNAKE_CASE						)

    for i, t in enumerate(scheduler.prk_timesteps						):
     UpperCAmelCase	:					Tuple	=	model(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)
     UpperCAmelCase	:					Optional[int]	=	scheduler.step_prk(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						).prev_sample

    for i, t in enumerate(scheduler.plms_timesteps						):
     UpperCAmelCase	:					List[Any]	=	model(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)
     UpperCAmelCase	:					Optional[int]	=	scheduler.step_plms(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						).prev_sample

    return sample

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      List[Any]:





    '''simple docstring'''




    UpperCAmelCase	:					str	=	dict(self.forward_default_kwargs						)

    UpperCAmelCase	:					Any	=	kwargs.pop("""num_inference_steps"""      ,     _SCREAMING_SNAKE_CASE						)

    for scheduler_class in self.scheduler_classes:
     UpperCAmelCase	:					str	=	self.get_scheduler_config()
     UpperCAmelCase	:					Tuple	=	scheduler_class(**_SCREAMING_SNAKE_CASE						)

     UpperCAmelCase	:					List[str]	=	self.dummy_sample
     UpperCAmelCase	:					str	=	0.1 * sample

     if num_inference_steps is not None and hasattr(_SCREAMING_SNAKE_CASE      ,     """set_timesteps"""						):
      scheduler.set_timesteps(_SCREAMING_SNAKE_CASE						)
     elif num_inference_steps is not None and not hasattr(_SCREAMING_SNAKE_CASE      ,     """set_timesteps"""						):
      UpperCAmelCase	:					str	=	num_inference_steps

     # copy over dummy past residuals (must be done after set_timesteps)
     UpperCAmelCase	:					Tuple	=	[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
     UpperCAmelCase	:					Any	=	dummy_past_residuals[:]

     UpperCAmelCase	:					Tuple	=	scheduler.step_prk(_SCREAMING_SNAKE_CASE      ,     0      ,     _SCREAMING_SNAKE_CASE      ,     **_SCREAMING_SNAKE_CASE						).prev_sample
     UpperCAmelCase	:					Optional[int]	=	scheduler.step_prk(_SCREAMING_SNAKE_CASE      ,     1      ,     _SCREAMING_SNAKE_CASE      ,     **_SCREAMING_SNAKE_CASE						).prev_sample

     self.assertEqual(output_a.shape      ,     sample.shape						)
     self.assertEqual(output_a.shape      ,     output_a.shape						)

     UpperCAmelCase	:					Dict	=	scheduler.step_plms(_SCREAMING_SNAKE_CASE      ,     0      ,     _SCREAMING_SNAKE_CASE      ,     **_SCREAMING_SNAKE_CASE						).prev_sample
     UpperCAmelCase	:					Tuple	=	scheduler.step_plms(_SCREAMING_SNAKE_CASE      ,     1      ,     _SCREAMING_SNAKE_CASE      ,     **_SCREAMING_SNAKE_CASE						).prev_sample

     self.assertEqual(output_a.shape      ,     sample.shape						)
     self.assertEqual(output_a.shape      ,     output_a.shape						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Any:





    '''simple docstring'''




    for timesteps in [100, 1000]:
     self.check_over_configs(num_train_timesteps=_SCREAMING_SNAKE_CASE						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      List[Any]:





    '''simple docstring'''




    for steps_offset in [0, 1]:
     self.check_over_configs(steps_offset=_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					int	=	self.scheduler_classes[0]
    UpperCAmelCase	:					int	=	self.get_scheduler_config(steps_offset=1						)
    UpperCAmelCase	:					Optional[int]	=	scheduler_class(**_SCREAMING_SNAKE_CASE						)
    scheduler.set_timesteps(10						)
    assert torch.equal(
        scheduler.timesteps      ,     torch.LongTensor(
            [901, 851, 851, 801, 801, 751, 751, 701, 701, 651, 651, 601, 601, 501, 401, 301, 201, 101, 1]						)      ,     )

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Optional[int]:





    '''simple docstring'''




    for beta_start, beta_end in zip([0.0001, 0.001]      ,     [0.002, 0.02]						):
     self.check_over_configs(beta_start=_SCREAMING_SNAKE_CASE      ,     beta_end=_SCREAMING_SNAKE_CASE						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Any:





    '''simple docstring'''




    for schedule in ["linear", "squaredcos_cap_v2"]:
     self.check_over_configs(beta_schedule=_SCREAMING_SNAKE_CASE						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Optional[int]:





    '''simple docstring'''




    for prediction_type in ["epsilon", "v_prediction"]:
     self.check_over_configs(prediction_type=_SCREAMING_SNAKE_CASE						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Optional[Any]:





    '''simple docstring'''




    for t in [1, 5, 10]:
     self.check_over_forward(time_step=_SCREAMING_SNAKE_CASE						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      List[str]:





    '''simple docstring'''




    for t, num_inference_steps in zip([1, 5, 10]      ,     [10, 50, 100]						):
     self.check_over_forward(num_inference_steps=_SCREAMING_SNAKE_CASE						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Optional[Any]:





    '''simple docstring'''




    UpperCAmelCase	:					Optional[Any]	=	27

    for scheduler_class in self.scheduler_classes:
     UpperCAmelCase	:					Tuple	=	self.dummy_sample
     UpperCAmelCase	:					str	=	0.1 * sample

     UpperCAmelCase	:					Dict	=	self.get_scheduler_config()
     UpperCAmelCase	:					Union[str, Any]	=	scheduler_class(**_SCREAMING_SNAKE_CASE						)

     scheduler.set_timesteps(_SCREAMING_SNAKE_CASE						)

     # before power of 3 fix, would error on first step, so we only need to do two
     for i, t in enumerate(scheduler.prk_timesteps[:2]						):
      UpperCAmelCase	:					Dict	=	scheduler.step_prk(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						).prev_sample

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      int:





    '''simple docstring'''




    with self.assertRaises(_SCREAMING_SNAKE_CASE						):
     UpperCAmelCase	:					Union[str, Any]	=	self.scheduler_classes[0]
     UpperCAmelCase	:					int	=	self.get_scheduler_config()
     UpperCAmelCase	:					Optional[int]	=	scheduler_class(**_SCREAMING_SNAKE_CASE						)

     scheduler.step_plms(self.dummy_sample      ,     1      ,     self.dummy_sample						).prev_sample

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      int:





    '''simple docstring'''




    UpperCAmelCase	:					int	=	self.full_loop()
    UpperCAmelCase	:					Optional[int]	=	torch.sum(torch.abs(_SCREAMING_SNAKE_CASE						)						)
    UpperCAmelCase	:					Union[str, Any]	=	torch.mean(torch.abs(_SCREAMING_SNAKE_CASE						)						)

    assert abs(result_sum.item() - 198.1318						) < 1E-2
    assert abs(result_mean.item() - 0.2580						) < 1E-3

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      List[Any]:





    '''simple docstring'''




    UpperCAmelCase	:					str	=	self.full_loop(prediction_type="""v_prediction"""						)
    UpperCAmelCase	:					Optional[Any]	=	torch.sum(torch.abs(_SCREAMING_SNAKE_CASE						)						)
    UpperCAmelCase	:					List[Any]	=	torch.mean(torch.abs(_SCREAMING_SNAKE_CASE						)						)

    assert abs(result_sum.item() - 67.3986						) < 1E-2
    assert abs(result_mean.item() - 0.0878						) < 1E-3

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Optional[Any]:





    '''simple docstring'''




    UpperCAmelCase	:					Optional[int]	=	self.full_loop(set_alpha_to_one=_SCREAMING_SNAKE_CASE      ,     beta_start=0.01						)
    UpperCAmelCase	:					Union[str, Any]	=	torch.sum(torch.abs(_SCREAMING_SNAKE_CASE						)						)
    UpperCAmelCase	:					List[str]	=	torch.mean(torch.abs(_SCREAMING_SNAKE_CASE						)						)

    assert abs(result_sum.item() - 230.0399						) < 1E-2
    assert abs(result_mean.item() - 0.2995						) < 1E-3




   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Optional[Any]:





    '''simple docstring'''




    UpperCAmelCase	:					str	=	self.full_loop(set_alpha_to_one=_SCREAMING_SNAKE_CASE      ,     beta_start=0.01						)
    UpperCAmelCase	:					Any	=	torch.sum(torch.abs(_SCREAMING_SNAKE_CASE						)						)
    UpperCAmelCase	:					str	=	torch.mean(torch.abs(_SCREAMING_SNAKE_CASE						)						)

    assert abs(result_sum.item() - 186.9482						) < 1E-2
    assert abs(result_mean.item() - 0.2434						) < 1E-3





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



@dataclass
class       SCREAMING_SNAKE_CASE__  (					datasets.BuilderConfig ):
   __lowerCAmelCase  :   Optional[datasets.Features]						=						None
   __lowerCAmelCase  :   str						=						"utf-8"
   __lowerCAmelCase  :   Optional[str]						=						None
   __lowerCAmelCase  :   Optional[str]						=						None
   __lowerCAmelCase  :   bool						=						True  # deprecated
   __lowerCAmelCase  :   Optional[int]						=						None  # deprecated
   __lowerCAmelCase  :   int						=						10 << 20  # 10MB
   __lowerCAmelCase  :   Optional[bool]						=						None



class       SCREAMING_SNAKE_CASE__  (					datasets.ArrowBasedBuilder ):
   __lowerCAmelCase  :   int						=						JsonConfig

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      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"""						)
     UpperCAmelCase	:					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 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE						)  ->      Tuple:





    '''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}"						)
    UpperCAmelCase	:					List[Any]	=	dl_manager.download_and_extract(self.config.data_files						)
    if isinstance(_SCREAMING_SNAKE_CASE      ,     (str, list, tuple)						):
     UpperCAmelCase	:					int	=	data_files
     if isinstance(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						):
      UpperCAmelCase	:					Tuple	=	[files]
     UpperCAmelCase	:					List[str]	=	[dl_manager.iter_files(_SCREAMING_SNAKE_CASE						) for file in files]
     return [datasets.SplitGenerator(name=datasets.Split.TRAIN      ,     gen_kwargs={"""files""": files}						)]
    UpperCAmelCase	:					List[str]	=	[]
    for split_name, files in data_files.items():
     if isinstance(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						):
      UpperCAmelCase	:					Dict	=	[files]
     UpperCAmelCase	:					str	=	[dl_manager.iter_files(_SCREAMING_SNAKE_CASE						) for file in files]
     splits.append(datasets.SplitGenerator(name=_SCREAMING_SNAKE_CASE      ,     gen_kwargs={"""files""": files}						)						)
    return splits

   def 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE						)  ->      pa.Table:





    '''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						):
      UpperCAmelCase	:					Tuple	=	self.config.features.arrow_schema.field(_SCREAMING_SNAKE_CASE						).type
      UpperCAmelCase	:					int	=	pa_table.append_column(_SCREAMING_SNAKE_CASE      ,     pa.array([None] * len(_SCREAMING_SNAKE_CASE						)      ,     type=_SCREAMING_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
     UpperCAmelCase	:					Any	=	table_cast(_SCREAMING_SNAKE_CASE      ,     self.config.features.arrow_schema						)
    return pa_table




   def 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE						)  ->      Optional[Any]:





    '''simple docstring'''




    for file_idx, file in enumerate(itertools.chain.from_iterable(_SCREAMING_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(_SCREAMING_SNAKE_CASE      ,     encoding=self.config.encoding      ,     errors=self.config.encoding_errors						) as f:
       UpperCAmelCase	:					Union[str, Any]	=	json.load(_SCREAMING_SNAKE_CASE						)

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

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

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