"""simple docstring"""





import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
  import jax
  import jax.numpy as jnp

  from transformers.generation import (
      FlaxForcedBOSTokenLogitsProcessor,
      FlaxForcedEOSTokenLogitsProcessor,
      FlaxLogitsProcessorList,
      FlaxMinLengthLogitsProcessor,
      FlaxTemperatureLogitsWarper,
      FlaxTopKLogitsWarper,
      FlaxTopPLogitsWarper,
  )
@require_flax
class 				lowercase__       (   unittest.TestCase				):







       '''simple docstring'''


       def 					lowercase__    (							self      :  Optional[int]		,    _UpperCAmelCase      :  Tuple		,    _UpperCAmelCase      :  Optional[Any]   )			->						int:


             '''simple docstring'''




             UpperCAmelCase_      			=							jnp.ones((batch_size, length)   ) / length
             return scores



       def 					lowercase__    (							self      :  List[Any]   )			->						Optional[int]:


             '''simple docstring'''




             UpperCAmelCase_      			=							None
             UpperCAmelCase_      			=							20

             UpperCAmelCase_      			=							self._get_uniform_logits(batch_size=2		,    length=lowerCamelCase_   )

             # tweak scores to not be uniform anymore
             UpperCAmelCase_      			=							scores.at[1, 5].set((1 / length) + 0.1   )  # peak, 1st batch
             UpperCAmelCase_      			=							scores.at[1, 10].set((1 / length) - 0.4   )  # valley, 1st batch

             # compute softmax
             UpperCAmelCase_      			=							jax.nn.softmax(lowerCamelCase_		,    axis=-1   )

             UpperCAmelCase_      			=							FlaxTemperatureLogitsWarper(temperature=0.5   )
             UpperCAmelCase_      			=							FlaxTemperatureLogitsWarper(temperature=1.3   )

             UpperCAmelCase_      			=							jax.nn.softmax(temp_dist_warper_sharper(lowerCamelCase_		,    scores.copy()		,    cur_len=lowerCamelCase_   )		,    axis=-1   )
             UpperCAmelCase_      			=							jax.nn.softmax(temp_dist_warper_smoother(lowerCamelCase_		,    scores.copy()		,    cur_len=lowerCamelCase_   )		,    axis=-1   )

             # uniform distribution stays uniform
             self.assertTrue(jnp.allclose(probs[0, :]		,    warped_prob_sharp[0, :]		,    atol=1e-3   )   )
             self.assertTrue(jnp.allclose(probs[0, :]		,    warped_prob_smooth[0, :]		,    atol=1e-3   )   )

             # sharp peaks get higher, valleys get lower
             self.assertLess(probs[1, :].max()		,    warped_prob_sharp[1, :].max()   )
             self.assertGreater(probs[1, :].min()		,    warped_prob_sharp[1, :].min()   )

             # smooth peaks get lower, valleys get higher
             self.assertGreater(probs[1, :].max()		,    warped_prob_smooth[1, :].max()   )
             self.assertLess(probs[1, :].min()		,    warped_prob_smooth[1, :].min()   )



       def 					lowercase__    (							self      :  Tuple   )			->						Any:


             '''simple docstring'''




             UpperCAmelCase_      			=							None
             UpperCAmelCase_      			=							10
             UpperCAmelCase_      			=							2

             # create ramp distribution
             UpperCAmelCase_      			=							np.broadcast_to(np.arange(lowerCamelCase_   )[None, :]		,    (batch_size, vocab_size)   ).copy()
             UpperCAmelCase_      			=							ramp_logits[1:, : vocab_size // 2] + vocab_size

             UpperCAmelCase_      			=							FlaxTopKLogitsWarper(3   )

             UpperCAmelCase_      			=							top_k_warp(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )

             # check that correct tokens are filtered
             self.assertListEqual(jnp.isinf(scores[0]   ).tolist()		,    7 * [True] + 3 * [False]   )
             self.assertListEqual(jnp.isinf(scores[1]   ).tolist()		,    2 * [True] + 3 * [False] + 5 * [True]   )

             # check special case
             UpperCAmelCase_      			=							5
             UpperCAmelCase_      			=							FlaxTopKLogitsWarper(top_k=1		,    filter_value=0.0		,    min_tokens_to_keep=3   )

             UpperCAmelCase_      			=							np.broadcast_to(np.arange(lowerCamelCase_   )[None, :]		,    (batch_size, length)   ).copy()
             UpperCAmelCase_      			=							top_k_warp_safety_check(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )

             # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
             self.assertListEqual((scores == 0.0).sum(axis=-1   ).tolist()		,    [2, 2]   )



       def 					lowercase__    (							self      :  str   )			->						Tuple:


             '''simple docstring'''




             UpperCAmelCase_      			=							None
             UpperCAmelCase_      			=							10
             UpperCAmelCase_      			=							2

             # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
             UpperCAmelCase_      			=							np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]   )   )

             UpperCAmelCase_      			=							FlaxTopPLogitsWarper(0.8   )
             UpperCAmelCase_      			=							np.exp(top_p_warp(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )   )

             # dist should be filtered to keep min num values so that sum is >= top_p
             # exp (-inf) => 0
             UpperCAmelCase_      			=							np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]   )
             self.assertTrue(np.allclose(lowerCamelCase_		,    lowerCamelCase_		,    atol=1e-3   )   )

             # check edge cases with negative and extreme logits
             UpperCAmelCase_      			=							np.broadcast_to(np.arange(lowerCamelCase_   )[None, :]		,    (batch_size, vocab_size)   ).copy() - (
                 vocab_size // 2
             )

             # make ramp_logits more extreme
             UpperCAmelCase_      			=							ramp_logits[1] * 100.0

             # make sure at least 2 tokens are kept
             UpperCAmelCase_      			=							FlaxTopPLogitsWarper(0.9		,    min_tokens_to_keep=2		,    filter_value=0.0   )
             UpperCAmelCase_      			=							top_p_warp(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )

             # first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
             self.assertListEqual((filtered_dist != 0.0).sum(axis=-1   ).tolist()		,    [3, 2]   )



       def 					lowercase__    (							self      :  Any   )			->						Dict:


             '''simple docstring'''




             UpperCAmelCase_      			=							20
             UpperCAmelCase_      			=							4
             UpperCAmelCase_      			=							0

             UpperCAmelCase_      			=							FlaxMinLengthLogitsProcessor(min_length=10		,    eos_token_id=lowerCamelCase_   )

             # check that min length is applied at length 5
             UpperCAmelCase_      			=							ids_tensor((batch_size, 20)		,    vocab_size=20   )
             UpperCAmelCase_      			=							5
             UpperCAmelCase_      			=							self._get_uniform_logits(lowerCamelCase_		,    lowerCamelCase_   )
             UpperCAmelCase_      			=							min_dist_processor(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
             self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()		,    4 * [-float("inf"   )]   )

             # check that min length is not applied anymore at length 15
             UpperCAmelCase_      			=							self._get_uniform_logits(lowerCamelCase_		,    lowerCamelCase_   )
             UpperCAmelCase_      			=							15
             UpperCAmelCase_      			=							min_dist_processor(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
             self.assertFalse(jnp.isinf(lowerCamelCase_   ).any()   )



       def 					lowercase__    (							self      :  List[str]   )			->						Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							20
             UpperCAmelCase_      			=							4
             UpperCAmelCase_      			=							0

             UpperCAmelCase_      			=							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase_   )

             # check that all scores are -inf except the bos_token_id score
             UpperCAmelCase_      			=							ids_tensor((batch_size, 1)		,    vocab_size=20   )
             UpperCAmelCase_      			=							1
             UpperCAmelCase_      			=							self._get_uniform_logits(lowerCamelCase_		,    lowerCamelCase_   )
             UpperCAmelCase_      			=							logits_processor(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
             self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]   ).all()   )
             self.assertListEqual(scores[:, bos_token_id].tolist()		,    4 * [0]   )  # score for bos_token_id shold be zero

             # check that bos_token_id is not forced if current length is greater than 1
             UpperCAmelCase_      			=							3
             UpperCAmelCase_      			=							self._get_uniform_logits(lowerCamelCase_		,    lowerCamelCase_   )
             UpperCAmelCase_      			=							logits_processor(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
             self.assertFalse(jnp.isinf(lowerCamelCase_   ).any()   )



       def 					lowercase__    (							self      :  Optional[Any]   )			->						Dict:


             '''simple docstring'''




             UpperCAmelCase_      			=							20
             UpperCAmelCase_      			=							4
             UpperCAmelCase_      			=							0
             UpperCAmelCase_      			=							5

             UpperCAmelCase_      			=							FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase_		,    eos_token_id=lowerCamelCase_   )

             # check that all scores are -inf except the eos_token_id when max_length is reached
             UpperCAmelCase_      			=							ids_tensor((batch_size, 4)		,    vocab_size=20   )
             UpperCAmelCase_      			=							4
             UpperCAmelCase_      			=							self._get_uniform_logits(lowerCamelCase_		,    lowerCamelCase_   )
             UpperCAmelCase_      			=							logits_processor(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
             self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]   ).all()   )
             self.assertListEqual(scores[:, eos_token_id].tolist()		,    4 * [0]   )  # score for eos_token_id should be zero

             # check that eos_token_id is not forced if max_length is not reached
             UpperCAmelCase_      			=							3
             UpperCAmelCase_      			=							self._get_uniform_logits(lowerCamelCase_		,    lowerCamelCase_   )
             UpperCAmelCase_      			=							logits_processor(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
             self.assertFalse(jnp.isinf(lowerCamelCase_   ).any()   )



       def 					lowercase__    (							self      :  int   )			->						List[Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							4
             UpperCAmelCase_      			=							10
             UpperCAmelCase_      			=							15
             UpperCAmelCase_      			=							2
             UpperCAmelCase_      			=							1
             UpperCAmelCase_      			=							15

             # dummy input_ids and scores
             UpperCAmelCase_      			=							ids_tensor((batch_size, sequence_length)		,    lowerCamelCase_   )
             UpperCAmelCase_      			=							input_ids.copy()

             UpperCAmelCase_      			=							self._get_uniform_logits(lowerCamelCase_		,    lowerCamelCase_   )
             UpperCAmelCase_      			=							scores.copy()

             # instantiate all dist processors
             UpperCAmelCase_      			=							FlaxTemperatureLogitsWarper(temperature=0.5   )
             UpperCAmelCase_      			=							FlaxTopKLogitsWarper(3   )
             UpperCAmelCase_      			=							FlaxTopPLogitsWarper(0.8   )

             # instantiate all logits processors
             UpperCAmelCase_      			=							FlaxMinLengthLogitsProcessor(min_length=10		,    eos_token_id=lowerCamelCase_   )
             UpperCAmelCase_      			=							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase_   )
             UpperCAmelCase_      			=							FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase_		,    eos_token_id=lowerCamelCase_   )

             UpperCAmelCase_      			=							10

             # no processor list
             UpperCAmelCase_      			=							temp_dist_warp(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
             UpperCAmelCase_      			=							top_k_warp(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
             UpperCAmelCase_      			=							top_p_warp(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
             UpperCAmelCase_      			=							min_dist_proc(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
             UpperCAmelCase_      			=							bos_dist_proc(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
             UpperCAmelCase_      			=							eos_dist_proc(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )

             # with processor list
             UpperCAmelCase_      			=							FlaxLogitsProcessorList(
                 [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]   )
             UpperCAmelCase_      			=							processor(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )

             # scores should be equal
             self.assertTrue(jnp.allclose(lowerCamelCase_		,    lowerCamelCase_		,    atol=1e-3   )   )

             # input_ids should never be changed
             self.assertListEqual(input_ids.tolist()		,    input_ids_comp.tolist()   )







       def 					lowercase__    (							self      :  Tuple   )			->						Tuple:


             '''simple docstring'''




             UpperCAmelCase_      			=							4
             UpperCAmelCase_      			=							10
             UpperCAmelCase_      			=							15
             UpperCAmelCase_      			=							2
             UpperCAmelCase_      			=							1
             UpperCAmelCase_      			=							15

             # dummy input_ids and scores
             UpperCAmelCase_      			=							ids_tensor((batch_size, sequence_length)		,    lowerCamelCase_   )
             UpperCAmelCase_      			=							input_ids.copy()

             UpperCAmelCase_      			=							self._get_uniform_logits(lowerCamelCase_		,    lowerCamelCase_   )
             UpperCAmelCase_      			=							scores.copy()

             # instantiate all dist processors
             UpperCAmelCase_      			=							FlaxTemperatureLogitsWarper(temperature=0.5   )
             UpperCAmelCase_      			=							FlaxTopKLogitsWarper(3   )
             UpperCAmelCase_      			=							FlaxTopPLogitsWarper(0.8   )

             # instantiate all logits processors
             UpperCAmelCase_      			=							FlaxMinLengthLogitsProcessor(min_length=10		,    eos_token_id=lowerCamelCase_   )
             UpperCAmelCase_      			=							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase_   )
             UpperCAmelCase_      			=							FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase_		,    eos_token_id=lowerCamelCase_   )

             UpperCAmelCase_      			=							10

             # no processor list
             def run_no_processor_list(_UpperCAmelCase      :  Optional[int]		,    _UpperCAmelCase      :  Optional[int]		,    _UpperCAmelCase      :  Optional[Any]   ):
                   UpperCAmelCase_      			=							temp_dist_warp(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
                   UpperCAmelCase_      			=							top_k_warp(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
                   UpperCAmelCase_      			=							top_p_warp(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
                   UpperCAmelCase_      			=							min_dist_proc(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
                   UpperCAmelCase_      			=							bos_dist_proc(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
                   UpperCAmelCase_      			=							eos_dist_proc(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
                   return scores

             # with processor list
             def run_processor_list(_UpperCAmelCase      :  Any		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  Any   ):
                   UpperCAmelCase_      			=							FlaxLogitsProcessorList(
                       [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]   )
                   UpperCAmelCase_      			=							processor(lowerCamelCase_		,    lowerCamelCase_		,    cur_len=lowerCamelCase_   )
                   return scores

             UpperCAmelCase_      			=							jax.jit(lowerCamelCase_   )
             UpperCAmelCase_      			=							jax.jit(lowerCamelCase_   )

             UpperCAmelCase_      			=							jitted_run_no_processor_list(lowerCamelCase_		,    lowerCamelCase_		,    lowerCamelCase_   )
             UpperCAmelCase_      			=							jitted_run_processor_list(lowerCamelCase_		,    lowerCamelCase_		,    lowerCamelCase_   )

             # scores should be equal
             self.assertTrue(jnp.allclose(lowerCamelCase_		,    lowerCamelCase_		,    atol=1e-3   )   )

             # input_ids should never be changed
             self.assertListEqual(input_ids.tolist()		,    input_ids_comp.tolist()   )









"""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       	=						datasets.utils.logging.get_logger(__name__)
@dataclass
class 				lowercase__       (   datasets.BuilderConfig				):







       '''simple docstring'''


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



class 				lowercase__       (   datasets.ArrowBasedBuilder				):







       '''simple docstring'''


       UpperCamelCase						=					JsonConfig



       def 					lowercase__    (							self      :  List[Any]   )			->						Dict:


             '''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_      			=							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 					lowercase__    (							self      :  Tuple		,    _UpperCAmelCase      :  Optional[int]   )			->						Optional[int]:


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



       def 					lowercase__    (							self      :  str		,    _UpperCAmelCase      :  pa.Table   )			->						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_      			=							self.config.features.arrow_schema.field(_UpperCAmelCase   ).type
                         UpperCAmelCase_      			=							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
                   UpperCAmelCase_      			=							table_cast(_UpperCAmelCase		,    self.config.features.arrow_schema   )
             return pa_table







       def 					lowercase__    (							self      :  Optional[int]		,    _UpperCAmelCase      :  List[Any]   )			->						str:


             '''simple docstring'''




             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:
                               UpperCAmelCase_      			=							json.load(_UpperCAmelCase   )

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

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




