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


UpperCAmelCase_								=							datasets.utils.logging.get_logger(__name__)





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





class      lowerCAmelCase						(					datasets.ArrowBasedBuilder			):
     _SCREAMING_SNAKE_CASE     :     int		=JsonConfig



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



     def 			a__  (					self				,      lowerCAmelCase__       ):
            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(lowerCAmelCase__       ).type
                          _A=	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
                   _A=	table_cast(lowerCAmelCase__				,      self.config.features.arrow_schema       )
            return pa_table





     def 			a__  (					self				,      lowerCAmelCase__       ):
            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:
                                 _A=	json.load(lowerCAmelCase__       )

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

from transformers import AutoConfig, AutoTokenizer, BertConfig, TensorType, is_flax_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, slow


if is_flax_available():
      import jax

      from transformers.models.auto.modeling_flax_auto import FlaxAutoModel
      from transformers.models.bert.modeling_flax_bert import FlaxBertModel
      from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel





@require_flax
class      lowerCAmelCase						(					unittest.TestCase			):



     @slow
     def 			a__  (					self       ):
            for model_name in ["bert-base-cased", "bert-large-uncased"]:
                   with self.subTest(lowerCAmelCase__       ):
                          _A=	AutoConfig.from_pretrained(lowerCAmelCase__       )
                          self.assertIsNotNone(lowerCAmelCase__       )
                          self.assertIsInstance(lowerCAmelCase__				,      lowerCAmelCase__       )

                          _A=	FlaxAutoModel.from_pretrained(lowerCAmelCase__       )
                          self.assertIsNotNone(lowerCAmelCase__       )
                          self.assertIsInstance(lowerCAmelCase__				,      lowerCAmelCase__       )



     @slow
     def 			a__  (					self       ):
            for model_name in ["roberta-base", "roberta-large"]:
                   with self.subTest(lowerCAmelCase__       ):
                          _A=	AutoConfig.from_pretrained(lowerCAmelCase__       )
                          self.assertIsNotNone(lowerCAmelCase__       )
                          self.assertIsInstance(lowerCAmelCase__				,      lowerCAmelCase__       )

                          _A=	FlaxAutoModel.from_pretrained(lowerCAmelCase__       )
                          self.assertIsNotNone(lowerCAmelCase__       )
                          self.assertIsInstance(lowerCAmelCase__				,      lowerCAmelCase__       )



     @slow
     def 			a__  (					self       ):
            for model_name in ["bert-base-cased", "bert-large-uncased"]:
                   _A=	AutoTokenizer.from_pretrained(lowerCAmelCase__       )
                   _A=	FlaxBertModel.from_pretrained(lowerCAmelCase__       )
                   _A=	tokenizer('Do you support jax jitted function?'				,      return_tensors=TensorType.JAX       )

                   @jax.jit
                   def eval(**lowerCAmelCase__       ):
                          return model(**lowerCAmelCase__       )

                   eval(**lowerCAmelCase__       ).block_until_ready()



     @slow
     def 			a__  (					self       ):
            for model_name in ["roberta-base", "roberta-large"]:
                   _A=	AutoTokenizer.from_pretrained(lowerCAmelCase__       )
                   _A=	FlaxRobertaModel.from_pretrained(lowerCAmelCase__       )
                   _A=	tokenizer('Do you support jax jitted function?'				,      return_tensors=TensorType.JAX       )

                   @jax.jit
                   def eval(**lowerCAmelCase__       ):
                          return model(**lowerCAmelCase__       )

                   eval(**lowerCAmelCase__       ).block_until_ready()



     def 			a__  (					self       ):
            with self.assertRaisesRegex(
                lowerCAmelCase__				,      'bert-base is not a local folder and is not a valid model identifier'       ):
                   _A=	FlaxAutoModel.from_pretrained('bert-base'       )



     def 			a__  (					self       ):
            with self.assertRaisesRegex(
                lowerCAmelCase__				,      r'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'       ):
                   _A=	FlaxAutoModel.from_pretrained(lowerCAmelCase__				,      revision='aaaaaa'       )



     def 			a__  (					self       ):
            with self.assertRaisesRegex(
                lowerCAmelCase__				,      'hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack'				,      ):
                   _A=	FlaxAutoModel.from_pretrained('hf-internal-testing/config-no-model'       )





     def 			a__  (					self       ):
            with self.assertRaisesRegex(lowerCAmelCase__				,      'Use `from_pt=True` to load this model'       ):
                   _A=	FlaxAutoModel.from_pretrained('hf-internal-testing/tiny-bert-pt-only'       )