




import inspect
import unittest
from typing import List

import numpy as np

from transformers import EfficientFormerConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import (
      TFEfficientFormerForImageClassification,
      TFEfficientFormerForImageClassificationWithTeacher,
      TFEfficientFormerModel,
  )
  from transformers.models.efficientformer.modeling_tf_efficientformer import (
      TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
  )


if is_vision_available():
  from PIL import Image

  from transformers import EfficientFormerImageProcessor
class 	_UpperCamelCase					:
      def __init__(	self      :Optional[int]       ,					lowerCamelCase      :List[Any]       ,					lowerCamelCase      :int = 13       ,					lowerCamelCase      :int = 64       ,					lowerCamelCase      :int = 2       ,					lowerCamelCase      :int = 3       ,					lowerCamelCase      :int = 3       ,					lowerCamelCase      :bool = True       ,					lowerCamelCase      :bool = True       ,					lowerCamelCase      :int = 128       ,					lowerCamelCase      :Optional[Any]=[16, 32, 64, 128]       ,					lowerCamelCase      :int = 7       ,					lowerCamelCase      :int = 4       ,					lowerCamelCase      :int = 37       ,					lowerCamelCase      :str = "gelu"       ,					lowerCamelCase      :float = 0.1       ,					lowerCamelCase      :float = 0.1       ,					lowerCamelCase      :int = 10       ,					lowerCamelCase      :float = 0.02       ,					lowerCamelCase      :int = 2       ,					lowerCamelCase      :int = 1       ,					lowerCamelCase      :int = 128       ,					lowerCamelCase      :List[int] = [2, 2, 2, 2]       ,					lowerCamelCase      :int = 2       ,					lowerCamelCase      :int = 2       ,					)   ->	List[Any]:
       UpperCAmelCase__							       =  parent
       UpperCAmelCase__							       =  batch_size
       UpperCAmelCase__							       =  image_size
       UpperCAmelCase__							       =  patch_size
       UpperCAmelCase__							       =  num_channels
       UpperCAmelCase__							       =  is_training
       UpperCAmelCase__							       =  use_labels
       UpperCAmelCase__							       =  hidden_size
       UpperCAmelCase__							       =  num_hidden_layers
       UpperCAmelCase__							       =  num_attention_heads
       UpperCAmelCase__							       =  intermediate_size
       UpperCAmelCase__							       =  hidden_act
       UpperCAmelCase__							       =  hidden_dropout_prob
       UpperCAmelCase__							       =  attention_probs_dropout_prob
       UpperCAmelCase__							       =  type_sequence_label_size
       UpperCAmelCase__							       =  initializer_range
       UpperCAmelCase__							       =  encoder_stride
       UpperCAmelCase__							       =  num_attention_outputs
       UpperCAmelCase__							       =  embed_dim
       UpperCAmelCase__							       =  embed_dim + 1
       UpperCAmelCase__							       =  resolution
       UpperCAmelCase__							       =  depths
       UpperCAmelCase__							       =  hidden_sizes
       UpperCAmelCase__							       =  dim
       UpperCAmelCase__							       =  mlp_expansion_ratio
      def 		UpperCAmelCase_  (	self      :List[Any]					)   ->	str:
       UpperCAmelCase__							       =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

       UpperCAmelCase__							       =  None
       if self.use_labels:
        UpperCAmelCase__							       =  ids_tensor([self.batch_size]       ,					self.type_sequence_label_size					)

       UpperCAmelCase__							       =  self.get_config()

       return config, pixel_values, labels
      def 		UpperCAmelCase_  (	self      :int					)   ->	str:
       return EfficientFormerConfig(
           image_size=self.image_size       ,					patch_size=self.patch_size       ,					num_channels=self.num_channels       ,					hidden_size=self.hidden_size       ,					num_hidden_layers=self.num_hidden_layers       ,					num_attention_heads=self.num_attention_heads       ,					intermediate_size=self.intermediate_size       ,					hidden_act=self.hidden_act       ,					hidden_dropout_prob=self.hidden_dropout_prob       ,					attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,					is_decoder=UpperCAmelCase__       ,					initializer_range=self.initializer_range       ,					encoder_stride=self.encoder_stride       ,					resolution=self.resolution       ,					depths=self.depths       ,					hidden_sizes=self.hidden_sizes       ,					dim=self.dim       ,					mlp_expansion_ratio=self.mlp_expansion_ratio       ,					)
      def 		UpperCAmelCase_  (	self      :Optional[int]       ,					lowerCamelCase      :Dict       ,					lowerCamelCase      :List[str]       ,					lowerCamelCase      :Dict					)   ->	Dict:
       UpperCAmelCase__							       =  TFEfficientFormerModel(config=UpperCAmelCase__					)
       UpperCAmelCase__							       =  model(UpperCAmelCase__       ,					training=UpperCAmelCase__					)
       self.parent.assertEqual(result.last_hidden_state.shape       ,					(self.batch_size, self.seq_length, self.hidden_size)					)
      def 		UpperCAmelCase_  (	self      :int       ,					lowerCamelCase      :Dict       ,					lowerCamelCase      :Dict       ,					lowerCamelCase      :str					)   ->	Union[str, Any]:
       UpperCAmelCase__							       =  self.type_sequence_label_size
       UpperCAmelCase__							       =  TFEfficientFormerForImageClassification(UpperCAmelCase__					)
       UpperCAmelCase__							       =  model(UpperCAmelCase__       ,					labels=UpperCAmelCase__       ,					training=UpperCAmelCase__					)
       self.parent.assertEqual(result.logits.shape       ,					(self.batch_size, self.type_sequence_label_size)					)

       # test greyscale images
       UpperCAmelCase__							       =  1
       UpperCAmelCase__							       =  TFEfficientFormerForImageClassification(UpperCAmelCase__					)

       UpperCAmelCase__							       =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
       UpperCAmelCase__							       =  model(UpperCAmelCase__       ,					labels=UpperCAmelCase__					)
       self.parent.assertEqual(result.logits.shape       ,					(self.batch_size, self.type_sequence_label_size)					)







      def 		UpperCAmelCase_  (	self      :int					)   ->	List[str]:
       UpperCAmelCase__							       =  self.prepare_config_and_inputs()
       UpperCAmelCase__       ,       UpperCAmelCase__       ,       UpperCAmelCase__							       =  config_and_inputs
       UpperCAmelCase__							       =  {"pixel_values": pixel_values}
       return config, inputs_dict



@require_tf
class 	_UpperCamelCase					( UpperCAmelCase__		,				UpperCAmelCase__		,				unittest.TestCase						):
      UpperCAmelCase_						=						(
          (
              TFEfficientFormerModel,
              TFEfficientFormerForImageClassificationWithTeacher,
              TFEfficientFormerForImageClassification,
          )
          if is_tf_available()
          else ()
      )
      UpperCAmelCase_						=						(
          {
              """feature-extraction""": TFEfficientFormerModel,
              """image-classification""": (
                  TFEfficientFormerForImageClassification,
                  TFEfficientFormerForImageClassificationWithTeacher,
              ),
          }
          if is_tf_available()
          else {}
      )

      UpperCAmelCase_						=						False

      UpperCAmelCase_						=						False
      UpperCAmelCase_						=						False
      UpperCAmelCase_						=						False
      UpperCAmelCase_						=						False
      def 		UpperCAmelCase_  (	self      :Optional[Any]					)   ->	List[str]:
       UpperCAmelCase__							       =  TFEfficientFormerModelTester(self					)
       UpperCAmelCase__							       =  ConfigTester(
           self       ,					config_class=UpperCAmelCase__       ,					has_text_modality=UpperCAmelCase__       ,					hidden_size=37					)
      def 		UpperCAmelCase_  (	self      :int					)   ->	Any:
       self.config_tester.run_common_tests()
      @unittest.skip(reason="EfficientFormer does not use inputs_embeds"					)
      def 		UpperCAmelCase_  (	self      :List[str]					)   ->	Dict:
       pass
      @unittest.skip(reason="EfficientFormer does not support input and output embeddings"					)
      def 		UpperCAmelCase_  (	self      :List[Any]					)   ->	Optional[Any]:
       pass
      def 		UpperCAmelCase_  (	self      :Any					)   ->	Optional[Any]:
       UpperCAmelCase__       ,       UpperCAmelCase__							       =  self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCAmelCase__							       =  model_class(UpperCAmelCase__					)
        UpperCAmelCase__							       =  inspect.signature(model.call					)
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        UpperCAmelCase__							       =  [*signature.parameters.keys()]

        UpperCAmelCase__							       =  ["pixel_values"]
        self.assertListEqual(arg_names[:1]       ,					UpperCAmelCase__					)
      def 		UpperCAmelCase_  (	self      :str					)   ->	Any:

       def check_hidden_states_output(lowerCamelCase      :Optional[Any]       ,					lowerCamelCase      :Any       ,					lowerCamelCase      :Dict					):
        UpperCAmelCase__							       =  model_class(UpperCAmelCase__					)

        UpperCAmelCase__							       =  model(**self._prepare_for_class(UpperCAmelCase__       ,					UpperCAmelCase__					)       ,					training=UpperCAmelCase__					)
        UpperCAmelCase__							       =  outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

        UpperCAmelCase__							       =  getattr(
            self.model_tester       ,					"expected_num_hidden_layers"       ,					self.model_tester.num_hidden_layers + 1					)
        self.assertEqual(len(UpperCAmelCase__					)       ,					UpperCAmelCase__					)

        if hasattr(self.model_tester       ,					"encoder_seq_length"					):
         UpperCAmelCase__							       =  self.model_tester.encoder_seq_length
         if hasattr(self.model_tester       ,					"chunk_length"					) and self.model_tester.chunk_length > 1:
          UpperCAmelCase__							       =  seq_length * self.model_tester.chunk_length
        else:
         UpperCAmelCase__							       =  self.model_tester.seq_length

        self.assertListEqual(
            list(hidden_states[-1].shape[-2:]					)       ,					[seq_length, self.model_tester.hidden_size]       ,					)

        if config.is_encoder_decoder:
         UpperCAmelCase__							       =  outputs.decoder_hidden_states

         self.asseretIsInstance(UpperCAmelCase__       ,					(list, tuple)					)
         self.assertEqual(len(UpperCAmelCase__					)       ,					UpperCAmelCase__					)
         UpperCAmelCase__							       =  getattr(self.model_tester       ,					"seq_length"       ,					UpperCAmelCase__					)
         UpperCAmelCase__							       =  getattr(self.model_tester       ,					"decoder_seq_length"       ,					UpperCAmelCase__					)

         self.assertListEqual(
             list(hidden_states[-1].shape[-2:]					)       ,					[decoder_seq_length, self.model_tester.hidden_size]       ,					)

       UpperCAmelCase__       ,       UpperCAmelCase__							       =  self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCAmelCase__							       =  True
        check_hidden_states_output(UpperCAmelCase__       ,					UpperCAmelCase__       ,					UpperCAmelCase__					)

        # check that output_hidden_states also work using config
        del inputs_dict["output_hidden_states"]
        UpperCAmelCase__							       =  True

        check_hidden_states_output(UpperCAmelCase__       ,					UpperCAmelCase__       ,					UpperCAmelCase__					)
      def 		UpperCAmelCase_  (	self      :Optional[Any]       ,					lowerCamelCase      :Dict       ,					lowerCamelCase      :Dict       ,					lowerCamelCase      :Dict=False					)   ->	int:
       UpperCAmelCase__							       =  super()._prepare_for_class(UpperCAmelCase__       ,					UpperCAmelCase__       ,					return_labels=UpperCAmelCase__					)

       if return_labels:
        if model_class.__name__ == "TFEfficientFormerForImageClassificationWithTeacher":
         del inputs_dict["labels"]

       return inputs_dict
      def 		UpperCAmelCase_  (	self      :Optional[Any]					)   ->	Union[str, Any]:
       UpperCAmelCase__							       =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*UpperCAmelCase__					)
      @unittest.skip(reason="EfficientFormer does not implement masked image modeling yet"					)
      def 		UpperCAmelCase_  (	self      :str					)   ->	str:
       UpperCAmelCase__							       =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_image_modeling(*UpperCAmelCase__					)
      def 		UpperCAmelCase_  (	self      :Any					)   ->	Tuple:
       UpperCAmelCase__							       =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase__					)
      @slow
      def 		UpperCAmelCase_  (	self      :Tuple					)   ->	Optional[int]:

       for model_name in TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCAmelCase__							       =  TFEfficientFormerModel.from_pretrained(UpperCAmelCase__					)
        self.assertIsNotNone(UpperCAmelCase__					)
      def 		UpperCAmelCase_  (	self      :Any					)   ->	str:
       UpperCAmelCase__       ,       UpperCAmelCase__							       =  self.model_tester.prepare_config_and_inputs_for_common()
       UpperCAmelCase__							       =  True

       UpperCAmelCase__							       =  getattr(self.model_tester       ,					"seq_length"       ,					UpperCAmelCase__					)
       UpperCAmelCase__							       =  getattr(self.model_tester       ,					"encoder_seq_length"       ,					UpperCAmelCase__					)
       UpperCAmelCase__							       =  getattr(self.model_tester       ,					"key_length"       ,					UpperCAmelCase__					)
       UpperCAmelCase__							       =  getattr(self.model_tester       ,					"chunk_length"       ,					UpperCAmelCase__					)

       if chunk_length is not None and hasattr(self.model_tester       ,					"num_hashes"					):
        UpperCAmelCase__							       =  encoder_seq_length * self.model_tester.num_hashes

       for model_class in self.all_model_classes:
        UpperCAmelCase__							       =  True
        UpperCAmelCase__							       =  False
        UpperCAmelCase__							       =  True
        UpperCAmelCase__							       =  model_class(UpperCAmelCase__					)

        UpperCAmelCase__							       =  model(**self._prepare_for_class(UpperCAmelCase__       ,					UpperCAmelCase__					)       ,					training=UpperCAmelCase__					)
        UpperCAmelCase__							       =  outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
        self.assertEqual(len(UpperCAmelCase__					)       ,					self.model_tester.num_attention_outputs					)

        # check that output_attentions also work using config
        del inputs_dict["output_attentions"]
        UpperCAmelCase__							       =  True
        UpperCAmelCase__							       =  model_class(UpperCAmelCase__					)
        UpperCAmelCase__							       =  model(**self._prepare_for_class(UpperCAmelCase__       ,					UpperCAmelCase__					)       ,					training=UpperCAmelCase__					)

        UpperCAmelCase__							       =  outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
        self.assertEqual(len(UpperCAmelCase__					)       ,					self.model_tester.num_attention_outputs					)

        if chunk_length is not None:
         self.assertListEqual(
             list(attentions[0].shape[-4:]					)       ,					[self.model_tester.num_attention_heads, encoder_seq_length, chunk_length, encoder_key_length]       ,					)
        else:
         self.assertListEqual(
             list(attentions[0].shape[-3:]					)       ,					[self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length]       ,					)







      def 		UpperCAmelCase_  (	self      :List[str]					)   ->	Optional[Any]:
       UpperCAmelCase__       ,       UpperCAmelCase__							       =  self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        # Prepare our model
        UpperCAmelCase__							       =  model_class(UpperCAmelCase__					)
        # These are maximally general inputs for the model, with multiple None dimensions
        # Hopefully this will catch any conditionals that fail for flexible shapes
        UpperCAmelCase__							       =  {
            key: tf.keras.Input(shape=val.shape[1:]       ,					dtype=val.dtype       ,					name=UpperCAmelCase__					)
            for key, val in model.input_signature.items()
            if key in model.dummy_inputs
        }
        UpperCAmelCase__							       =  model(UpperCAmelCase__					)
        self.assertTrue(outputs_dict is not None					)



def   lowerCAmelCase						(						):




 """simple docstring"""







 UpperCAmelCase__							       =  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"		)
 return image


@require_tf
@require_vision
class 	_UpperCamelCase					( unittest.TestCase						):
      @cached_property
      def 		UpperCAmelCase_  (	self      :List[str]					)   ->	List[str]:
       return (
           EfficientFormerImageProcessor.from_pretrained("snap-research/efficientformer-l1-300"					)
           if is_vision_available()
           else None
       )
      @slow
      def 		UpperCAmelCase_  (	self      :List[str]					)   ->	Any:
       UpperCAmelCase__							       =  TFEfficientFormerForImageClassification.from_pretrained("snap-research/efficientformer-l1-300"					)
       UpperCAmelCase__							       =  self.default_image_processor
       UpperCAmelCase__							       =  prepare_img()
       UpperCAmelCase__							       =  image_processor(images=UpperCAmelCase__       ,					return_tensors="tf"					)
       # forward pass
       UpperCAmelCase__							       =  model(**UpperCAmelCase__       ,					training=UpperCAmelCase__					)
       # verify the logits
       UpperCAmelCase__							       =  tf.TensorShape((1, 1000)					)
       self.assertEqual(outputs.logits.shape       ,					UpperCAmelCase__					)
       UpperCAmelCase__							       =  tf.constant([-0.05_55, 0.48_25, -0.08_52]					)
       self.assertTrue(np.allclose(outputs.logits[0, :3]       ,					UpperCAmelCase__       ,					atol=1e-4					)					)







      @slow
      def 		UpperCAmelCase_  (	self      :Dict					)   ->	int:
       UpperCAmelCase__							       =  TFEfficientFormerForImageClassificationWithTeacher.from_pretrained(
           "snap-research/efficientformer-l1-300"					)
       UpperCAmelCase__							       =  self.default_image_processor
       UpperCAmelCase__							       =  prepare_img()
       UpperCAmelCase__							       =  image_processor(images=UpperCAmelCase__       ,					return_tensors="tf"					)
       # forward pass
       UpperCAmelCase__							       =  model(**UpperCAmelCase__       ,					training=UpperCAmelCase__					)
       # verify the logits
       UpperCAmelCase__							       =  tf.TensorShape((1, 1000)					)
       self.assertEqual(outputs.logits.shape       ,					UpperCAmelCase__					)
       UpperCAmelCase__							       =  tf.constant([-0.13_12, 0.43_53, -1.04_99]					)
       self.assertTrue(np.allclose(outputs.logits[0, :3]       ,					UpperCAmelCase__       ,					atol=1e-4					)					)











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


@dataclass
class 				UpperCamelCase_    ( 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 				UpperCamelCase_    ( datasets.ArrowBasedBuilder     ):

    '''simple docstring'''






    UpperCAmelCase__							  =      JsonConfig


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





           '''simple docstring'''

           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      :	List[str]					, UpperCAmelCase__      :	List[Any])						->Dict:





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


    def        SCREAMING_SNAKE_CASE							(			self      :	Optional[Any]					, 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):
                         A__       =    self.config.features.arrow_schema.field(UpperCAmelCase__).type
                         A__       =    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
                  A__       =    table_cast(UpperCAmelCase__					, self.config.features.arrow_schema)
           return pa_table





    def        SCREAMING_SNAKE_CASE							(			self      :	Union[str, Any]					, UpperCAmelCase__      :	Tuple)						->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:
                                A__       =    json.load(UpperCAmelCase__)

                         # 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(UpperCAmelCase__					, (list, tuple)):
                                A__       =    set().union(*[row.keys() for row in dataset])
                                A__       =    {col: [row.get(UpperCAmelCase__) for row in dataset] for col in keys}
                         else:
                                A__       =    dataset
                         A__       =    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:
                                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(UpperCAmelCase__)
                                       # PyArrow only accepts utf-8 encoded bytes
                                       if self.config.encoding != "utf-8":
                                              A__       =    batch.decode(self.config.encoding					, errors=UpperCAmelCase__).encode('''utf-8''')
                                       try:
                                              while True:
                                                     try:
                                                            A__       =    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:
                                                            A__       =    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:
                                                            A__       =    set().union(*[row.keys() for row in dataset])
                                                            A__       =    {col: [row.get(UpperCAmelCase__) for row in dataset] for col in keys}
                                                            A__       =    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


