'''simple docstring'''

from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFSequenceClassifierOutput,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import logging
from .configuration_regnet import RegNetConfig


lowerCAmelCase_ : int						  =      logging.get_logger(__name__)

# General docstring
lowerCAmelCase_ : List[Any]						  =      "RegNetConfig"

# Base docstring
lowerCAmelCase_ : List[str]						  =      "facebook/regnet-y-040"
lowerCAmelCase_ : List[Any]						  =      [1, 1_088, 7, 7]

# Image classification docstring
lowerCAmelCase_ : Tuple						  =      "facebook/regnet-y-040"
lowerCAmelCase_ : Optional[Any]						  =      "tabby, tabby cat"

lowerCAmelCase_ : int						  =      [
    "facebook/regnet-y-040",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]





class 							lowercase						(						tf.keras.layers.Layer ):


   def __init__(   self				:	Union[str, Any]   ,		__lowerCAmelCase				:	int   ,		__lowerCAmelCase				:	int = 3   ,		__lowerCAmelCase				:	int = 1   ,		__lowerCAmelCase				:	int = 1   ,		__lowerCAmelCase				:	Optional[str] = "relu"   ,		**__lowerCAmelCase				:	List[str]   ,		)      ->					Tuple:
         super().__init__(**__lowerCAmelCase)
         # The padding and conv has been verified in
         # https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
         lowercase_     =      tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2)
         lowercase_     =      tf.keras.layers.ConvaD(
             filters=__lowerCAmelCase   ,		kernel_size=__lowerCAmelCase   ,		strides=__lowerCAmelCase   ,		padding="VALID"   ,		groups=__lowerCAmelCase   ,		use_bias=__lowerCAmelCase   ,		name="convolution"   ,		)
         lowercase_     =      tf.keras.layers.BatchNormalization(epsilon=1e-5   ,		momentum=0.9   ,		name="normalization")
         lowercase_     =      ACTaFN[activation] if activation is not None else tf.identity


   def __UpperCAmelCase       (   self				:	Dict   ,		__lowerCAmelCase				:	List[Any])      ->					Optional[int]:
         lowercase_     =      self.convolution(self.padding(__lowerCAmelCase))
         lowercase_     =      self.normalization(__lowerCAmelCase)
         lowercase_     =      self.activation(__lowerCAmelCase)
         return hidden_state





class 							lowercase						(						tf.keras.layers.Layer ):


   def __init__(   self				:	Any   ,		__lowerCAmelCase				:	RegNetConfig   ,		**__lowerCAmelCase				:	str)      ->					Tuple:
         super().__init__(**__lowerCAmelCase)
         lowercase_     =      config.num_channels
         lowercase_     =      TFRegNetConvLayer(
             out_channels=config.embedding_size   ,		kernel_size=3   ,		stride=2   ,		activation=config.hidden_act   ,		name="embedder"   ,		)


   def __UpperCAmelCase       (   self				:	Optional[Any]   ,		__lowerCAmelCase				:	List[str])      ->					Tuple:
         lowercase_     =      shape_list(__lowerCAmelCase)[1]
         if tf.executing_eagerly() and num_channels != self.num_channels:
               raise ValueError(
                   "Make sure that the channel dimension of the pixel values match with the one set in the configuration.")

         # When running on CPU, `tf.keras.layers.Conv2D` doesn't support `NCHW` format.
         # So change the input format from `NCHW` to `NHWC`.
         # shape = (batch_size, in_height, in_width, in_channels=num_channels)
         lowercase_     =      tf.transpose(__lowerCAmelCase   ,		perm=(0, 2, 3, 1))
         lowercase_     =      self.embedder(__lowerCAmelCase)
         return hidden_state





class 							lowercase						(						tf.keras.layers.Layer ):


   def __init__(   self				:	Optional[Any]   ,		__lowerCAmelCase				:	int   ,		__lowerCAmelCase				:	int = 2   ,		**__lowerCAmelCase				:	Union[str, Any])      ->					List[str]:
         super().__init__(**__lowerCAmelCase)
         lowercase_     =      tf.keras.layers.ConvaD(
             filters=__lowerCAmelCase   ,		kernel_size=1   ,		strides=__lowerCAmelCase   ,		use_bias=__lowerCAmelCase   ,		name="convolution")
         lowercase_     =      tf.keras.layers.BatchNormalization(epsilon=1e-5   ,		momentum=0.9   ,		name="normalization")


   def __UpperCAmelCase       (   self				:	Any   ,		__lowerCAmelCase				:	tf.Tensor   ,		__lowerCAmelCase				:	bool = False)      ->					tf.Tensor:
         return self.normalization(self.convolution(__lowerCAmelCase)   ,		training=__lowerCAmelCase)





class 							lowercase						(						tf.keras.layers.Layer ):


   def __init__(   self				:	Optional[int]   ,		__lowerCAmelCase				:	int   ,		__lowerCAmelCase				:	int   ,		**__lowerCAmelCase				:	List[str])      ->					str:
         super().__init__(**__lowerCAmelCase)
         lowercase_     =      tf.keras.layers.GlobalAveragePoolingaD(keepdims=__lowerCAmelCase   ,		name="pooler")
         lowercase_     =      [
             tf.keras.layers.ConvaD(filters=__lowerCAmelCase   ,		kernel_size=1   ,		activation="relu"   ,		name="attention.0"),
             tf.keras.layers.ConvaD(filters=__lowerCAmelCase   ,		kernel_size=1   ,		activation="sigmoid"   ,		name="attention.2"),
         ]


   def __UpperCAmelCase       (   self				:	Optional[Any]   ,		__lowerCAmelCase				:	str)      ->					Dict:
         # [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
         lowercase_     =      self.pooler(__lowerCAmelCase)
         for layer_module in self.attention:
               lowercase_     =      layer_module(__lowerCAmelCase)
         lowercase_     =      hidden_state * pooled
         return hidden_state





class 							lowercase						(						tf.keras.layers.Layer ):


   def __init__(   self				:	Tuple   ,		__lowerCAmelCase				:	RegNetConfig   ,		__lowerCAmelCase				:	int   ,		__lowerCAmelCase				:	int   ,		__lowerCAmelCase				:	int = 1   ,		**__lowerCAmelCase				:	Optional[Any])      ->					Dict:
         super().__init__(**__lowerCAmelCase)
         lowercase_     =      in_channels != out_channels or stride != 1
         lowercase_     =      max(1   ,		out_channels // config.groups_width)
         lowercase_     =      (
             TFRegNetShortCut(__lowerCAmelCase   ,		stride=__lowerCAmelCase   ,		name="shortcut")
             if should_apply_shortcut
             else tf.keras.layers.Activation("linear"   ,		name="shortcut")
         )
         # `self.layers` instead of `self.layer` because that is a reserved argument.
         lowercase_     =      [
             TFRegNetConvLayer(__lowerCAmelCase   ,		kernel_size=1   ,		activation=config.hidden_act   ,		name="layer.0"),
             TFRegNetConvLayer(
                 __lowerCAmelCase   ,		stride=__lowerCAmelCase   ,		groups=__lowerCAmelCase   ,		activation=config.hidden_act   ,		name="layer.1"),
             TFRegNetConvLayer(__lowerCAmelCase   ,		kernel_size=1   ,		activation=__lowerCAmelCase   ,		name="layer.2"),
         ]
         lowercase_     =      ACTaFN[config.hidden_act]


   def __UpperCAmelCase       (   self				:	str   ,		__lowerCAmelCase				:	int)      ->					Tuple:
         lowercase_     =      hidden_state
         for layer_module in self.layers:
               lowercase_     =      layer_module(__lowerCAmelCase)
         lowercase_     =      self.shortcut(__lowerCAmelCase)
         hidden_state += residual
         lowercase_     =      self.activation(__lowerCAmelCase)
         return hidden_state





class 							lowercase						(						tf.keras.layers.Layer ):


   def __init__(   self				:	Optional[int]   ,		__lowerCAmelCase				:	RegNetConfig   ,		__lowerCAmelCase				:	int   ,		__lowerCAmelCase				:	int   ,		__lowerCAmelCase				:	int = 1   ,		**__lowerCAmelCase				:	int)      ->					Tuple:
         super().__init__(**__lowerCAmelCase)
         lowercase_     =      in_channels != out_channels or stride != 1
         lowercase_     =      max(1   ,		out_channels // config.groups_width)
         lowercase_     =      (
             TFRegNetShortCut(__lowerCAmelCase   ,		stride=__lowerCAmelCase   ,		name="shortcut")
             if should_apply_shortcut
             else tf.keras.layers.Activation("linear"   ,		name="shortcut")
         )
         lowercase_     =      [
             TFRegNetConvLayer(__lowerCAmelCase   ,		kernel_size=1   ,		activation=config.hidden_act   ,		name="layer.0"),
             TFRegNetConvLayer(
                 __lowerCAmelCase   ,		stride=__lowerCAmelCase   ,		groups=__lowerCAmelCase   ,		activation=config.hidden_act   ,		name="layer.1"),
             TFRegNetSELayer(__lowerCAmelCase   ,		reduced_channels=int(round(in_channels / 4))   ,		name="layer.2"),
             TFRegNetConvLayer(__lowerCAmelCase   ,		kernel_size=1   ,		activation=__lowerCAmelCase   ,		name="layer.3"),
         ]
         lowercase_     =      ACTaFN[config.hidden_act]


   def __UpperCAmelCase       (   self				:	Dict   ,		__lowerCAmelCase				:	List[Any])      ->					str:
         lowercase_     =      hidden_state
         for layer_module in self.layers:
               lowercase_     =      layer_module(__lowerCAmelCase)
         lowercase_     =      self.shortcut(__lowerCAmelCase)
         hidden_state += residual
         lowercase_     =      self.activation(__lowerCAmelCase)
         return hidden_state





class 							lowercase						(						tf.keras.layers.Layer ):


   def __init__(   self				:	List[str]   ,		__lowerCAmelCase				:	RegNetConfig   ,		__lowerCAmelCase				:	int   ,		__lowerCAmelCase				:	int   ,		__lowerCAmelCase				:	int = 2   ,		__lowerCAmelCase				:	int = 2   ,		**__lowerCAmelCase				:	str)      ->					int:
         super().__init__(**__lowerCAmelCase)

         lowercase_     =      TFRegNetXLayer if config.layer_type == "x" else TFRegNetYLayer
         lowercase_     =      [
             # downsampling is done in the first layer with stride of 2
             layer(__lowerCAmelCase   ,		__lowerCAmelCase   ,		__lowerCAmelCase   ,		stride=__lowerCAmelCase   ,		name="layers.0"),
             *[layer(__lowerCAmelCase   ,		__lowerCAmelCase   ,		__lowerCAmelCase   ,		name=F'layers.{i+1}') for i in range(depth - 1)],
         ]


   def __UpperCAmelCase       (   self				:	Union[str, Any]   ,		__lowerCAmelCase				:	Tuple)      ->					List[Any]:
         for layer_module in self.layers:
               lowercase_     =      layer_module(__lowerCAmelCase)
         return hidden_state





class 							lowercase						(						tf.keras.layers.Layer ):


   def __init__(   self				:	Optional[int]   ,		__lowerCAmelCase				:	RegNetConfig   ,		**__lowerCAmelCase				:	Union[str, Any])      ->					Optional[int]:
         super().__init__(**__lowerCAmelCase)
         lowercase_     =      []
         # based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
         self.stages.append(
             TFRegNetStage(
                 __lowerCAmelCase   ,		config.embedding_size   ,		config.hidden_sizes[0]   ,		stride=2 if config.downsample_in_first_stage else 1   ,		depth=config.depths[0]   ,		name="stages.0"   ,		))
         lowercase_     =      zip(config.hidden_sizes   ,		config.hidden_sizes[1:])
         for i, ((in_channels, out_channels), depth) in enumerate(zip(__lowerCAmelCase   ,		config.depths[1:])):
               self.stages.append(TFRegNetStage(__lowerCAmelCase   ,		__lowerCAmelCase   ,		__lowerCAmelCase   ,		depth=__lowerCAmelCase   ,		name=F'stages.{i+1}'))


   def __UpperCAmelCase       (   self				:	List[Any]   ,		__lowerCAmelCase				:	tf.Tensor   ,		__lowerCAmelCase				:	bool = False   ,		__lowerCAmelCase				:	bool = True)      ->					TFBaseModelOutputWithNoAttention:
         lowercase_     =      () if output_hidden_states else None

         for stage_module in self.stages:
               if output_hidden_states:
                     lowercase_     =      hidden_states + (hidden_state,)

               lowercase_     =      stage_module(__lowerCAmelCase)

         if output_hidden_states:
               lowercase_     =      hidden_states + (hidden_state,)

         if not return_dict:
               return tuple(v for v in [hidden_state, hidden_states] if v is not None)

         return TFBaseModelOutputWithNoAttention(last_hidden_state=__lowerCAmelCase   ,		hidden_states=__lowerCAmelCase)





@keras_serializable
class 							lowercase						(						tf.keras.layers.Layer ):
   lowerCamelCase_    					=RegNetConfig


   def __init__(   self				:	Any   ,		__lowerCAmelCase				:	Tuple   ,		**__lowerCAmelCase				:	Optional[int])      ->					Any:
         super().__init__(**__lowerCAmelCase)
         lowercase_     =      config
         lowercase_     =      TFRegNetEmbeddings(__lowerCAmelCase   ,		name="embedder")
         lowercase_     =      TFRegNetEncoder(__lowerCAmelCase   ,		name="encoder")
         lowercase_     =      tf.keras.layers.GlobalAveragePoolingaD(keepdims=__lowerCAmelCase   ,		name="pooler")


   @unpack_inputs
   def __UpperCAmelCase       (   self				:	Optional[int]   ,		__lowerCAmelCase				:	tf.Tensor   ,		__lowerCAmelCase				:	Optional[bool] = None   ,		__lowerCAmelCase				:	Optional[bool] = None   ,		__lowerCAmelCase				:	bool = False   ,		)      ->					TFBaseModelOutputWithPoolingAndNoAttention:
         lowercase_     =      (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         lowercase_     =      return_dict if return_dict is not None else self.config.use_return_dict

         lowercase_     =      self.embedder(__lowerCAmelCase   ,		training=__lowerCAmelCase)

         lowercase_     =      self.encoder(
             __lowerCAmelCase   ,		output_hidden_states=__lowerCAmelCase   ,		return_dict=__lowerCAmelCase   ,		training=__lowerCAmelCase)

         lowercase_     =      encoder_outputs[0]
         lowercase_     =      self.pooler(__lowerCAmelCase)

         # Change to NCHW output format have uniformity in the modules
         lowercase_     =      tf.transpose(__lowerCAmelCase   ,		perm=(0, 3, 1, 2))
         lowercase_     =      tf.transpose(__lowerCAmelCase   ,		perm=(0, 3, 1, 2))

         # Change the other hidden state outputs to NCHW as well
         if output_hidden_states:
               lowercase_     =      tuple([tf.transpose(__lowerCAmelCase   ,		perm=(0, 3, 1, 2)) for h in encoder_outputs[1]])

         if not return_dict:
               return (last_hidden_state, pooled_output) + encoder_outputs[1:]

         return TFBaseModelOutputWithPoolingAndNoAttention(
             last_hidden_state=__lowerCAmelCase   ,		pooler_output=__lowerCAmelCase   ,		hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states   ,		)





class 							lowercase						(						__lowerCamelCase ):
   lowerCamelCase_    					=RegNetConfig
   lowerCamelCase_    					='regnet'
   lowerCamelCase_    					='pixel_values'


   @property
   def __UpperCAmelCase       (   self				:	Dict)      ->					Union[str, Any]:
         return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)   ,		dtype=tf.floataa)}




lowerCAmelCase_ : Union[str, Any]						  =      r"\n    Parameters:\n    This model is a Tensorflow\n    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a\n    regular Tensorflow Module and refer to the Tensorflow documentation for all matter related to general usage and\n    behavior.\n        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.\n"

lowerCAmelCase_ : Optional[Any]						  =      r"\n    Args:\n        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConveNextImageProcessor.__call__`] for details.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"





@add_start_docstrings(
    'The bare RegNet model outputting raw features without any specific head on top.' ,				__lowerCamelCase ,				)
class 							lowercase						(						__lowerCamelCase ):


   def __init__(   self				:	str   ,		__lowerCAmelCase				:	RegNetConfig   ,		*__lowerCAmelCase				:	int   ,		**__lowerCAmelCase				:	Union[str, Any])      ->					Optional[int]:
         super().__init__(__lowerCAmelCase   ,		*__lowerCAmelCase   ,		**__lowerCAmelCase)
         lowercase_     =      TFRegNetMainLayer(__lowerCAmelCase   ,		name="regnet")


   @unpack_inputs
   @add_start_docstrings_to_model_forward(__lowerCAmelCase)
   @add_code_sample_docstrings(
       checkpoint=_CHECKPOINT_FOR_DOC   ,		output_type=__lowerCAmelCase   ,		config_class=_CONFIG_FOR_DOC   ,		modality="vision"   ,		expected_output=_EXPECTED_OUTPUT_SHAPE   ,		)
   def __UpperCAmelCase       (   self				:	List[Any]   ,		__lowerCAmelCase				:	tf.Tensor   ,		__lowerCAmelCase				:	Optional[bool] = None   ,		__lowerCAmelCase				:	Optional[bool] = None   ,		__lowerCAmelCase				:	Optional[int]=False   ,		)      ->					Union[TFBaseModelOutputWithPoolingAndNoAttention, Tuple[tf.Tensor]]:
         lowercase_     =      (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         lowercase_     =      return_dict if return_dict is not None else self.config.use_return_dict

         lowercase_     =      self.regnet(
             pixel_values=__lowerCAmelCase   ,		output_hidden_states=__lowerCAmelCase   ,		return_dict=__lowerCAmelCase   ,		training=__lowerCAmelCase   ,		)
         if not return_dict:
               return (outputs[0],) + outputs[1:]

         return TFBaseModelOutputWithPoolingAndNoAttention(
             last_hidden_state=outputs.last_hidden_state   ,		pooler_output=outputs.pooler_output   ,		hidden_states=outputs.hidden_states   ,		)





@add_start_docstrings(
    '\n    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    ' ,				__lowerCamelCase ,				)
class 							lowercase						(						__lowerCamelCase ,				__lowerCamelCase ):


   def __init__(   self				:	Any   ,		__lowerCAmelCase				:	RegNetConfig   ,		*__lowerCAmelCase				:	Any   ,		**__lowerCAmelCase				:	List[Any])      ->					Optional[Any]:
         super().__init__(__lowerCAmelCase   ,		*__lowerCAmelCase   ,		**__lowerCAmelCase)
         lowercase_     =      config.num_labels
         lowercase_     =      TFRegNetMainLayer(__lowerCAmelCase   ,		name="regnet")
         # classification head
         lowercase_     =      [
             tf.keras.layers.Flatten(),
             tf.keras.layers.Dense(config.num_labels   ,		name="classifier.1") if config.num_labels > 0 else tf.identity,
         ]


   @unpack_inputs
   @add_start_docstrings_to_model_forward(__lowerCAmelCase)
   @add_code_sample_docstrings(
       checkpoint=_IMAGE_CLASS_CHECKPOINT   ,		output_type=__lowerCAmelCase   ,		config_class=_CONFIG_FOR_DOC   ,		expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,		)
   def __UpperCAmelCase       (   self				:	List[Any]   ,		__lowerCAmelCase				:	tf.Tensor = None   ,		__lowerCAmelCase				:	tf.Tensor = None   ,		__lowerCAmelCase				:	bool = None   ,		__lowerCAmelCase				:	bool = None   ,		__lowerCAmelCase				:	Optional[int]=False   ,		)      ->					Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
         lowercase_     =      (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         lowercase_     =      return_dict if return_dict is not None else self.config.use_return_dict

         lowercase_     =      self.regnet(
             __lowerCAmelCase   ,		output_hidden_states=__lowerCAmelCase   ,		return_dict=__lowerCAmelCase   ,		training=__lowerCAmelCase)

         lowercase_     =      outputs.pooler_output if return_dict else outputs[1]

         lowercase_     =      self.classifier[0](__lowerCAmelCase)
         lowercase_     =      self.classifier[1](__lowerCAmelCase)

         lowercase_     =      None if labels is None else self.hf_compute_loss(labels=__lowerCAmelCase   ,		logits=__lowerCAmelCase)

         if not return_dict:
               lowercase_     =      (logits,) + outputs[2:]
               return ((loss,) + output) if loss is not None else output

         return TFSequenceClassifierOutput(loss=__lowerCAmelCase   ,		logits=__lowerCAmelCase   ,		hidden_states=outputs.hidden_states)







'''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_ : Union[str, Any]						  =      datasets.utils.logging.get_logger(__name__)





@dataclass
class 							lowercase						(						datasets.BuilderConfig ):
   lowerCamelCase_    					=None
   lowerCamelCase_    					="utf-8"
   lowerCamelCase_    					=None
   lowerCamelCase_    					=None
   lowerCamelCase_    					=True  # deprecated
   lowerCamelCase_    					=None  # deprecated
   lowerCamelCase_    					=1_0 << 2_0  # 10MB
   lowerCamelCase_    					=None





class 							lowercase						(						datasets.ArrowBasedBuilder ):
   lowerCamelCase_    					=JsonConfig


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


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


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

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

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

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






