





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 					lowercase__ (			datasets.BuilderConfig       ):

   '''simple docstring'''


   a		:							Optional[datasets.Features]     	=	None
   a		:							str     	=	"utf-8"
   a		:							Optional[str]     	=	None
   a		:							Optional[str]     	=	None
   a		:							bool     	=	True  # deprecated
   a		:							Optional[int]     	=	None  # deprecated
   a		:							int     	=	10 << 20  # 10MB
   a		:							Optional[bool]     	=	None






class 					lowercase__ (			datasets.ArrowBasedBuilder       ):

   '''simple docstring'''


   a		:							int     	=	JsonConfig




   def    UpperCamelCase__						(					self							)				->   Optional[Any]:






       """simple docstring"""
       if self.config.block_size is not None:
           logger.warning('''The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead'''							)
           UpperCamelCase__							:       str      =   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,	__magic_name__							)				->   Union[str, Any]:






       """simple docstring"""
       if not self.config.data_files:
           raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}"							)
       UpperCamelCase__							:       int      =   dl_manager.download_and_extract(self.config.data_files							)
       if isinstance(__magic_name__,	(str, list, tuple)							):
           UpperCamelCase__							:       List[str]      =   data_files
           if isinstance(__magic_name__,	__magic_name__							):
               UpperCamelCase__							:       Optional[Any]      =   [files]
           UpperCamelCase__							:       Tuple      =   [dl_manager.iter_files(__magic_name__							) for file in files]
           return [datasets.SplitGenerator(name=datasets.Split.TRAIN,	gen_kwargs={'''files''': files}							)]
       UpperCamelCase__							:       Union[str, Any]      =   []
       for split_name, files in data_files.items():
           if isinstance(__magic_name__,	__magic_name__							):
               UpperCamelCase__							:       Tuple      =   [files]
           UpperCamelCase__							:       List[str]      =   [dl_manager.iter_files(__magic_name__							) for file in files]
           splits.append(datasets.SplitGenerator(name=__magic_name__,	gen_kwargs={'''files''': files}							)							)
       return splits




   def    UpperCamelCase__						(					self,	__magic_name__							)				->   pa.Table:






       """simple docstring"""
       if self.config.features is not None:
           # adding missing columns
           for column_name in set(self.config.features							) - set(pa_table.column_names							):
               UpperCamelCase__							:       Any      =   self.config.features.arrow_schema.field(__magic_name__							).type
               UpperCamelCase__							:       Tuple      =   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
           UpperCamelCase__							:       int      =   table_cast(__magic_name__,	self.config.features.arrow_schema							)
       return pa_table




   def    UpperCamelCase__						(					self,	__magic_name__							)				->   Optional[int]:






       """simple docstring"""
       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:
                   UpperCamelCase__							:       Dict      =   json.load(__magic_name__							)

               # We keep only the field we are interested in
               UpperCamelCase__							:       Union[str, Any]      =   dataset[self.config.field]

               # We accept two format: a list of dicts or a dict of lists
               if isinstance(__magic_name__,	(list, tuple)							):
                   UpperCamelCase__							:       Any      =   set().union(*[row.keys() for row in dataset]							)
                   UpperCamelCase__							:       Optional[int]      =   {col: [row.get(__magic_name__							) for row in dataset] for col in keys}
               else:
                   UpperCamelCase__							:       int      =   dataset
               UpperCamelCase__							:       Union[str, Any]      =   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:
                   UpperCamelCase__							:       Union[str, Any]      =   0
                   # Use block_size equal to the chunk size divided by 32 to leverage multithreading
                   # Set a default minimum value of 16kB if the chunk size is really small
                   UpperCamelCase__							:       Union[str, Any]      =   max(self.config.chunksize // 32,	16 << 10							)
                   UpperCamelCase__							:       Optional[Any]      =   (
                       self.config.encoding_errors if self.config.encoding_errors is not None else '''strict'''
                   )
                   while True:
                       UpperCamelCase__							:       Union[str, Any]      =   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":
                           UpperCamelCase__							:       List[Any]      =   batch.decode(self.config.encoding,	errors=__magic_name__							).encode('''utf-8'''							)
                       try:
                           while True:
                               try:
                                   UpperCamelCase__							:       Union[str, Any]      =   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:
                                   UpperCamelCase__							:       List[str]      =   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:
                                   UpperCamelCase__							:       str      =   set().union(*[row.keys() for row in dataset]							)
                                   UpperCamelCase__							:       List[Any]      =   {col: [row.get(__magic_name__							) for row in dataset] for col in keys}
                                   UpperCamelCase__							:       Optional[Any]      =   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








import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
  import jax

  from transformers.models.big_bird.modeling_flax_big_bird import (
      FlaxBigBirdForCausalLM,
      FlaxBigBirdForMaskedLM,
      FlaxBigBirdForMultipleChoice,
      FlaxBigBirdForPreTraining,
      FlaxBigBirdForQuestionAnswering,
      FlaxBigBirdForSequenceClassification,
      FlaxBigBirdForTokenClassification,
      FlaxBigBirdModel,
  )






class 					lowercase__ (			unittest.TestCase       ):

   '''simple docstring'''


   def __init__(					self,	__magic_name__,	__magic_name__=2,	__magic_name__=56,	__magic_name__=True,	__magic_name__=True,	__magic_name__=True,	__magic_name__=True,	__magic_name__=99,	__magic_name__=32,	__magic_name__=2,	__magic_name__=2,	__magic_name__=7,	__magic_name__="gelu_new",	__magic_name__=0.1,	__magic_name__=0.1,	__magic_name__=512,	__magic_name__=16,	__magic_name__=2,	__magic_name__=0.02,	__magic_name__=4,	__magic_name__="block_sparse",	__magic_name__=True,	__magic_name__=False,	__magic_name__=2,	__magic_name__=3,	)				->   Optional[int]:






       """simple docstring"""
       UpperCamelCase__							:       List[str]      =   parent
       UpperCamelCase__							:       List[Any]      =   batch_size
       UpperCamelCase__							:       str      =   seq_length
       UpperCamelCase__							:       List[str]      =   is_training
       UpperCamelCase__							:       List[Any]      =   use_attention_mask
       UpperCamelCase__							:       Any      =   use_token_type_ids
       UpperCamelCase__							:       Tuple      =   use_labels
       UpperCamelCase__							:       Any      =   vocab_size
       UpperCamelCase__							:       Optional[int]      =   hidden_size
       UpperCamelCase__							:       Union[str, Any]      =   num_hidden_layers
       UpperCamelCase__							:       Optional[Any]      =   num_attention_heads
       UpperCamelCase__							:       List[str]      =   intermediate_size
       UpperCamelCase__							:       List[str]      =   hidden_act
       UpperCamelCase__							:       Union[str, Any]      =   hidden_dropout_prob
       UpperCamelCase__							:       Union[str, Any]      =   attention_probs_dropout_prob
       UpperCamelCase__							:       int      =   max_position_embeddings
       UpperCamelCase__							:       List[Any]      =   type_vocab_size
       UpperCamelCase__							:       Tuple      =   type_sequence_label_size
       UpperCamelCase__							:       Optional[Any]      =   initializer_range
       UpperCamelCase__							:       Dict      =   num_choices

       UpperCamelCase__							:       Dict      =   rescale_embeddings
       UpperCamelCase__							:       Tuple      =   attention_type
       UpperCamelCase__							:       Union[str, Any]      =   use_bias
       UpperCamelCase__							:       Union[str, Any]      =   block_size
       UpperCamelCase__							:       Optional[int]      =   num_random_blocks




   def    UpperCamelCase__						(					self							)				->   int:






       """simple docstring"""
       UpperCamelCase__							:       int      =   ids_tensor([self.batch_size, self.seq_length],	self.vocab_size							)

       UpperCamelCase__							:       int      =   None
       if self.use_attention_mask:
           UpperCamelCase__							:       List[str]      =   random_attention_mask([self.batch_size, self.seq_length]							)

       UpperCamelCase__							:       Optional[int]      =   None
       if self.use_token_type_ids:
           UpperCamelCase__							:       List[str]      =   ids_tensor([self.batch_size, self.seq_length],	self.type_vocab_size							)

       UpperCamelCase__							:       str      =   BigBirdConfig(
           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=__magic_name__,	initializer_range=self.initializer_range,	attention_type=self.attention_type,	block_size=self.block_size,	num_random_blocks=self.num_random_blocks,	use_bias=self.use_bias,	rescale_embeddings=self.rescale_embeddings,	)

       return config, input_ids, token_type_ids, attention_mask




   def    UpperCamelCase__						(					self							)				->   str:






       """simple docstring"""
       UpperCamelCase__							:       List[str]      =   self.prepare_config_and_inputs()
       UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__							:       Tuple      =   config_and_inputs
       UpperCamelCase__							:       Union[str, Any]      =   {
           '''input_ids''': input_ids,
           '''token_type_ids''': token_type_ids,
           '''attention_mask''': attention_mask,
       }
       return config, inputs_dict






@require_flax
class 					lowercase__ (			__lowerCamelCase    ,       unittest.TestCase       ):

   '''simple docstring'''


   a		:							Any     	=	(
       (
           FlaxBigBirdForCausalLM,
           FlaxBigBirdModel,
           FlaxBigBirdForPreTraining,
           FlaxBigBirdForMaskedLM,
           FlaxBigBirdForMultipleChoice,
           FlaxBigBirdForQuestionAnswering,
           FlaxBigBirdForSequenceClassification,
           FlaxBigBirdForTokenClassification,
       )
       if is_flax_available()
       else ()
   )

   a		:							List[str]     	=	False
   a		:							str     	=	False




   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__							:       Union[str, Any]      =   FlaxBigBirdModelTester(self							)




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Dict:






       """simple docstring"""
       super().test_from_pretrained_save_pretrained()




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Optional[Any]:






       """simple docstring"""
       super().test_from_pretrained_with_no_automatic_init()




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Dict:






       """simple docstring"""
       super().test_no_automatic_init()




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Optional[Any]:






       """simple docstring"""
       super().test_hidden_states_output()




   @slow
   def    UpperCamelCase__						(					self							)				->   List[str]:






       """simple docstring"""
       for model_class_name in self.all_model_classes:
           UpperCamelCase__							:       List[str]      =   model_class_name.from_pretrained('''google/bigbird-roberta-base'''							)
           self.assertIsNotNone(__magic_name__							)




   def    UpperCamelCase__						(					self							)				->   str:






       """simple docstring"""
       if self.test_attn_probs:
           super().test_attention_outputs()




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__     ,UpperCamelCase__							:       Dict      =   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
           with self.subTest(model_class.__name__							):
               UpperCamelCase__							:       Union[str, Any]      =   self._prepare_for_class(__magic_name__,	__magic_name__							)
               UpperCamelCase__							:       Tuple      =   model_class(__magic_name__							)

               @jax.jit
               def model_jitted(__magic_name__,	__magic_name__=None,	**__magic_name__							):
                   return model(input_ids=__magic_name__,	attention_mask=__magic_name__,	**__magic_name__							)

               with self.subTest('''JIT Enabled'''							):
                   UpperCamelCase__							:       str      =   model_jitted(**__magic_name__							).to_tuple()

               with self.subTest('''JIT Disabled'''							):
                   with jax.disable_jit():
                       UpperCamelCase__							:       Dict      =   model_jitted(**__magic_name__							).to_tuple()

               self.assertEqual(len(__magic_name__							),	len(__magic_name__							)							)
               for jitted_output, output in zip(__magic_name__,	__magic_name__							):
                   self.assertEqual(jitted_output.shape,	output.shape							)




   def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__,	__magic_name__,	__magic_name__=1E-5,	__magic_name__="outputs",	__magic_name__=None							)				->   List[Any]:






       """simple docstring"""
       # `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
       # an effort was done to return `attention_probs` (yet to be verified).
       if name.startswith('''outputs.attentions'''							):
           return
       else:
           super().check_pt_flax_outputs(__magic_name__,	__magic_name__,	__magic_name__,	__magic_name__,	__magic_name__,	__magic_name__							)

