

'''simple docstring'''






import unittest

from parameterized import parameterized

from transformers import AutoTokenizer, GPTNeoXConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import (
     GPTNeoXForCausalLM,
     GPTNeoXForQuestionAnswering,
     GPTNeoXForSequenceClassification,
     GPTNeoXForTokenClassification,
     GPTNeoXModel,
 )






class 			lowercase_	:




  """simple docstring"""




  def __init__(      self   :     Tuple					,lowercase__   :     int					,lowercase__   :     str=1_3					,lowercase__   :     Union[str, Any]=7					,lowercase__   :     List[str]=True					,lowercase__   :     List[str]=True					,lowercase__   :     List[str]=True					,lowercase__   :     Any=True					,lowercase__   :     str=9_9					,lowercase__   :     List[Any]=6_4					,lowercase__   :     Tuple=5					,lowercase__   :     List[Any]=4					,lowercase__   :     Any=3_7					,lowercase__   :     str="gelu"					,lowercase__   :     Any=0.1					,lowercase__   :     Union[str, Any]=0.1					,lowercase__   :     List[Any]=5_1_2					,lowercase__   :     int=1_6					,lowercase__   :     Dict=2					,lowercase__   :     Dict=0.0_2					,lowercase__   :     str=3					,lowercase__   :     Optional[Any]=4					,lowercase__   :     int=None					,):
   __lowercase							  = parent
   __lowercase							  = batch_size
   __lowercase							  = seq_length
   __lowercase							  = is_training
   __lowercase							  = use_input_mask
   __lowercase							  = use_token_type_ids
   __lowercase							  = use_labels
   __lowercase							  = vocab_size
   __lowercase							  = hidden_size
   __lowercase							  = num_hidden_layers
   __lowercase							  = num_attention_heads
   __lowercase							  = intermediate_size
   __lowercase							  = hidden_act
   __lowercase							  = hidden_dropout_prob
   __lowercase							  = attention_probs_dropout_prob
   __lowercase							  = max_position_embeddings
   __lowercase							  = type_vocab_size
   __lowercase							  = type_sequence_label_size
   __lowercase							  = initializer_range
   __lowercase							  = num_labels
   __lowercase							  = num_choices
   __lowercase							  = scope
   __lowercase							  = vocab_size - 1






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]			):
   __lowercase							  = ids_tensor([self.batch_size, self.seq_length]					,self.vocab_size			)

   __lowercase							  = None
   if self.use_input_mask:
    __lowercase							  = random_attention_mask([self.batch_size, self.seq_length]			)

   __lowercase							  = None
   if self.use_labels:
    __lowercase							  = ids_tensor([self.batch_size, self.seq_length]					,self.num_labels			)

   __lowercase							  = self.get_config()

   return config, input_ids, input_mask, token_labels






  def 					SCREAMING_SNAKE_CASE      (      self   :     int			):
   return GPTNeoXConfig(
       vocab_size=self.vocab_size					,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					,max_position_embeddings=self.max_position_embeddings					,type_vocab_size=self.type_vocab_size					,is_decoder=lowercase__					,initializer_range=self.initializer_range					,pad_token_id=self.pad_token_id					,)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   __lowercase  ,	__lowercase  ,	__lowercase  ,	__lowercase							  = self.prepare_config_and_inputs()

   __lowercase							  = True

   return config, input_ids, input_mask, token_labels






  def 					SCREAMING_SNAKE_CASE      (      self   :     str					,lowercase__   :     Tuple					,lowercase__   :     Dict					,lowercase__   :     Any			):
   __lowercase							  = GPTNeoXModel(config=lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(lowercase__					,attention_mask=lowercase__			)
   __lowercase							  = model(lowercase__			)
   self.parent.assertEqual(result.last_hidden_state.shape					,(self.batch_size, self.seq_length, self.hidden_size)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]					,lowercase__   :     Any					,lowercase__   :     Union[str, Any]					,lowercase__   :     int			):
   __lowercase							  = True
   __lowercase							  = GPTNeoXModel(lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(lowercase__					,attention_mask=lowercase__			)
   self.parent.assertEqual(result.last_hidden_state.shape					,(self.batch_size, self.seq_length, self.hidden_size)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]					,lowercase__   :     List[Any]					,lowercase__   :     Dict					,lowercase__   :     Optional[int]					,lowercase__   :     Dict			):
   __lowercase							  = GPTNeoXForCausalLM(config=lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(lowercase__					,attention_mask=lowercase__					,labels=lowercase__			)
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.seq_length, self.vocab_size)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]					,lowercase__   :     Optional[Any]					,lowercase__   :     Union[str, Any]					,lowercase__   :     Optional[Any]					,lowercase__   :     Optional[Any]			):
   __lowercase							  = self.num_labels
   __lowercase							  = GPTNeoXForQuestionAnswering(lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(lowercase__					,attention_mask=lowercase__			)
   self.parent.assertEqual(result.start_logits.shape					,(self.batch_size, self.seq_length)			)
   self.parent.assertEqual(result.end_logits.shape					,(self.batch_size, self.seq_length)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple					,lowercase__   :     List[Any]					,lowercase__   :     Dict					,lowercase__   :     List[str]					,lowercase__   :     str			):
   __lowercase							  = self.num_labels
   __lowercase							  = GPTNeoXForSequenceClassification(lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = ids_tensor([self.batch_size]					,self.type_sequence_label_size			)
   __lowercase							  = model(lowercase__					,attention_mask=lowercase__					,labels=lowercase__			)
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.num_labels)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Dict					,lowercase__   :     Any					,lowercase__   :     str					,lowercase__   :     Union[str, Any]					,lowercase__   :     Optional[int]			):
   __lowercase							  = self.num_labels
   __lowercase							  = GPTNeoXForTokenClassification(lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(lowercase__					,attention_mask=lowercase__					,labels=lowercase__			)
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.seq_length, self.num_labels)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]					,lowercase__   :     str					,lowercase__   :     Optional[Any]					,lowercase__   :     Optional[int]			):
   __lowercase							  = True
   __lowercase							  = GPTNeoXForCausalLM(config=lowercase__			)
   model.to(lowercase__			)
   model.eval()

   # first forward pass
   __lowercase							  = model(lowercase__					,attention_mask=lowercase__					,use_cache=lowercase__			)
   __lowercase							  = outputs.past_key_values

   # create hypothetical multiple next token and extent to next_input_ids
   __lowercase							  = ids_tensor((self.batch_size, 3)					,config.vocab_size			)
   __lowercase							  = ids_tensor((self.batch_size, 3)					,vocab_size=2			)

   # append to next input_ids and
   __lowercase							  = torch.cat([input_ids, next_tokens]					,dim=-1			)
   __lowercase							  = torch.cat([input_mask, next_mask]					,dim=-1			)

   __lowercase							  = model(lowercase__					,attention_mask=lowercase__					,output_hidden_states=lowercase__			)
   __lowercase							  = output_from_no_past['''hidden_states'''][0]
   __lowercase							  = model(
       lowercase__					,attention_mask=lowercase__					,past_key_values=lowercase__					,output_hidden_states=lowercase__					,)['''hidden_states'''][0]

   # select random slice
   __lowercase							  = ids_tensor((1,)					,output_from_past.shape[-1]			).item()
   __lowercase							  = output_from_no_past[:, -3:, random_slice_idx].detach()
   __lowercase							  = output_from_past[:, :, random_slice_idx].detach()

   self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]			)

   # test that outputs are equal for slice
   self.parent.assertTrue(torch.allclose(lowercase__					,lowercase__					,atol=1e-3			)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[Any]			):
   __lowercase							  = self.prepare_config_and_inputs()
   __lowercase  ,	__lowercase  ,	__lowercase  ,	__lowercase							  = config_and_inputs
   __lowercase							  = {'''input_ids''': input_ids, '''attention_mask''': input_mask}
   return config, inputs_dict






@require_torch
class 			lowercase_	(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       unittest.TestCase      ):




  """simple docstring"""




  SCREAMING_SNAKE_CASE    :  Optional[Any]							=			(
      (
          GPTNeoXModel,
          GPTNeoXForCausalLM,
          GPTNeoXForQuestionAnswering,
          GPTNeoXForSequenceClassification,
          GPTNeoXForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  SCREAMING_SNAKE_CASE    :  Union[str, Any]							=			(GPTNeoXForCausalLM,) if is_torch_available() else ()
  SCREAMING_SNAKE_CASE    :  Union[str, Any]							=			(
      {
          'feature-extraction': GPTNeoXModel,
          'question-answering': GPTNeoXForQuestionAnswering,
          'text-classification': GPTNeoXForSequenceClassification,
          'text-generation': GPTNeoXForCausalLM,
          'token-classification': GPTNeoXForTokenClassification,
          'zero-shot': GPTNeoXForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  SCREAMING_SNAKE_CASE    :  List[str]							=			False
  SCREAMING_SNAKE_CASE    :  Optional[int]							=			False
  SCREAMING_SNAKE_CASE    :  List[str]							=			False
  SCREAMING_SNAKE_CASE    :  Dict							=			False






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]			):
   __lowercase							  = GPTNeoXModelTester(self			)
   __lowercase							  = ConfigTester(self					,config_class=lowercase__					,hidden_size=6_4					,num_attention_heads=8			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]			):
   self.config_tester.run_common_tests()






  def 					SCREAMING_SNAKE_CASE      (      self   :     Any			):
   __lowercase  ,	__lowercase  ,	__lowercase  ,	__lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(lowercase__					,lowercase__					,lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     str			):
   __lowercase  ,	__lowercase  ,	__lowercase  ,	__lowercase							  = self.model_tester.prepare_config_and_inputs_for_decoder()
   self.model_tester.create_and_check_model_as_decoder(lowercase__					,lowercase__					,lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   # This regression test was failing with PyTorch < 1.3
   __lowercase  ,	__lowercase  ,	__lowercase  ,	__lowercase							  = self.model_tester.prepare_config_and_inputs_for_decoder()

   __lowercase							  = None

   self.model_tester.create_and_check_model_as_decoder(lowercase__					,lowercase__					,lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[Any]			):
   __lowercase  ,	__lowercase  ,	__lowercase  ,	__lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_decoder_model_past_large_inputs(lowercase__					,lowercase__					,lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_causal_lm(*lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_question_answering(*lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     str			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_sequence_classification(*lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_token_classification(*lowercase__			)






  @unittest.skip(reason='''Feed forward chunking is not implemented'''			)
  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   pass






  @parameterized.expand([('''linear''',), ('''dynamic''',)]			)
  def 					SCREAMING_SNAKE_CASE      (      self   :     str					,lowercase__   :     Union[str, Any]			):
   __lowercase  ,	__lowercase							  = self.model_tester.prepare_config_and_inputs_for_common()
   __lowercase							  = ids_tensor([1, 1_0]					,config.vocab_size			)
   __lowercase							  = ids_tensor([1, int(config.max_position_embeddings * 1.5			)]					,config.vocab_size			)

   set_seed(4_2			)  # Fixed seed at init time so the two models get the same random weights
   __lowercase							  = GPTNeoXModel(lowercase__			)
   original_model.to(lowercase__			)
   original_model.eval()
   __lowercase							  = original_model(lowercase__			).last_hidden_state
   __lowercase							  = original_model(lowercase__			).last_hidden_state

   set_seed(4_2			)  # Fixed seed at init time so the two models get the same random weights
   __lowercase							  = {'''type''': scaling_type, '''factor''': 1_0.0}
   __lowercase							  = GPTNeoXModel(lowercase__			)
   scaled_model.to(lowercase__			)
   scaled_model.eval()
   __lowercase							  = scaled_model(lowercase__			).last_hidden_state
   __lowercase							  = scaled_model(lowercase__			).last_hidden_state

   # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
   # maximum sequence length, so the outputs for the short input should match.
   if scaling_type == "dynamic":
    self.assertTrue(torch.allclose(lowercase__					,lowercase__					,atol=1e-5			)			)
   else:
    self.assertFalse(torch.allclose(lowercase__					,lowercase__					,atol=1e-5			)			)

   # The output should be different for long inputs
   self.assertFalse(torch.allclose(lowercase__					,lowercase__					,atol=1e-5			)			)






@require_torch
class 			lowercase_	(unittest.TestCase      ):




  """simple docstring"""




  @slow
  def 					SCREAMING_SNAKE_CASE      (      self   :     List[Any]			):
   __lowercase							  = AutoTokenizer.from_pretrained('''EleutherAI/pythia-410m-deduped'''			)
   for checkpointing in [True, False]:
    __lowercase							  = GPTNeoXForCausalLM.from_pretrained('''EleutherAI/pythia-410m-deduped'''			)

    if checkpointing:
     model.gradient_checkpointing_enable()
    else:
     model.gradient_checkpointing_disable()
    model.to(lowercase__			)

    __lowercase							  = tokenizer('''My favorite food is'''					,return_tensors='''pt'''			).to(lowercase__			)
    # The hub repo. is updated on 2023-04-04, resulting in poor outputs.
    # See: https://github.com/huggingface/transformers/pull/24193
    __lowercase							  = '''My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI\'m not sure'''

    __lowercase							  = model.generate(**lowercase__					,do_sample=lowercase__					,max_new_tokens=2_0			)
    __lowercase							  = tokenizer.batch_decode(lowercase__			)[0]

    self.assertEqual(lowercase__					,lowercase__			)


import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


A   :     Tuple       =				datasets.utils.logging.get_logger(__name__)


@dataclass
class      lowerCamelCase    (datasets.BuilderConfig    ):
  """simple docstring"""







  lowerCamelCase__   =  None
  lowerCamelCase__   =  "utf-8"
  lowerCamelCase__   =  None
  lowerCamelCase__   =  None
  lowerCamelCase__   =  True  # deprecated
  lowerCamelCase__   =  None  # deprecated
  lowerCamelCase__   =  1_0 << 2_0  # 10MB
  lowerCamelCase__   =  None



class      lowerCamelCase    (datasets.ArrowBasedBuilder    ):
  """simple docstring"""







  lowerCamelCase__   =  JsonConfig







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







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







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

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

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

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


