
import unittest

import numpy as np

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

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
							from transformers.models.bert.modeling_flax_bert import (
							    FlaxBertForMaskedLM,
							    FlaxBertForMultipleChoice,
							    FlaxBertForNextSentencePrediction,
							    FlaxBertForPreTraining,
							    FlaxBertForQuestionAnswering,
							    FlaxBertForSequenceClassification,
							    FlaxBertForTokenClassification,
							    FlaxBertModel,
							)

class     _UpperCamelCase			(   unittest.TestCase					):






			'''simple docstring'''

			def __init__(    self	:    List[str]  ,		snake_case_	:    str  ,		snake_case_	:    Dict=13  ,		snake_case_	:    List[str]=7  ,		snake_case_	:    Dict=True  ,		snake_case_	:    Union[str, Any]=True  ,		snake_case_	:    Union[str, Any]=True  ,		snake_case_	:    Dict=True  ,		snake_case_	:    Tuple=99  ,		snake_case_	:    str=32  ,		snake_case_	:    List[Any]=5  ,		snake_case_	:    str=4  ,		snake_case_	:    Dict=37  ,		snake_case_	:    Optional[int]="gelu"  ,		snake_case_	:    List[str]=0.1  ,		snake_case_	:    str=0.1  ,		snake_case_	:    Any=512  ,		snake_case_	:    Tuple=16  ,		snake_case_	:    Optional[int]=2  ,		snake_case_	:    str=0.02  ,		snake_case_	:    Tuple=4  ,		):
							UpperCamelCase_:							Optional[int]		=  parent
							UpperCamelCase_:							Tuple		=  batch_size
							UpperCamelCase_:							Optional[Any]		=  seq_length
							UpperCamelCase_:							List[str]		=  is_training
							UpperCamelCase_:							Optional[Any]		=  use_attention_mask
							UpperCamelCase_:							Any		=  use_token_type_ids
							UpperCamelCase_:							List[str]		=  use_labels
							UpperCamelCase_:							Tuple		=  vocab_size
							UpperCamelCase_:							List[str]		=  hidden_size
							UpperCamelCase_:							List[str]		=  num_hidden_layers
							UpperCamelCase_:							Dict		=  num_attention_heads
							UpperCamelCase_:							Any		=  intermediate_size
							UpperCamelCase_:							Any		=  hidden_act
							UpperCamelCase_:							List[Any]		=  hidden_dropout_prob
							UpperCamelCase_:							Dict		=  attention_probs_dropout_prob
							UpperCamelCase_:							int		=  max_position_embeddings
							UpperCamelCase_:							int		=  type_vocab_size
							UpperCamelCase_:							Dict		=  type_sequence_label_size
							UpperCamelCase_:							List[str]		=  initializer_range
							UpperCamelCase_:							List[str]		=  num_choices

			def  lowerCAmelCase__		(    self	:    Any ):
							UpperCamelCase_:							str		=  ids_tensor([self.batch_size, self.seq_length]  ,		self.vocab_size )

							UpperCamelCase_:							List[str]		=  None
							if self.use_attention_mask:
											UpperCamelCase_:							Optional[Any]		=  random_attention_mask([self.batch_size, self.seq_length] )

							UpperCamelCase_:							Dict		=  None
							if self.use_token_type_ids:
											UpperCamelCase_:							str		=  ids_tensor([self.batch_size, self.seq_length]  ,		self.type_vocab_size )

							UpperCamelCase_:							Any		=  BertConfig(
							    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=snake_case_  ,		initializer_range=self.initializer_range  ,		)

							return config, input_ids, token_type_ids, attention_mask

			def  lowerCAmelCase__		(    self	:    Optional[Any] ):
							UpperCamelCase_:							Optional[Any]		=  self.prepare_config_and_inputs()
							UpperCamelCase_,				UpperCamelCase_,				UpperCamelCase_,				UpperCamelCase_:							Dict		=  config_and_inputs
							UpperCamelCase_:							str		=  {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
							return config, inputs_dict




			def  lowerCAmelCase__		(    self	:    Any ):
							UpperCamelCase_:							Tuple		=  self.prepare_config_and_inputs()
							UpperCamelCase_,				UpperCamelCase_,				UpperCamelCase_,				UpperCamelCase_:							str		=  config_and_inputs

							UpperCamelCase_:							Dict		=  True
							UpperCamelCase_:							Optional[int]		=  floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
							UpperCamelCase_:							Any		=  ids_tensor([self.batch_size, self.seq_length]  ,		vocab_size=2 )

							return (
							    config,
							    input_ids,
							    attention_mask,
							    encoder_hidden_states,
							    encoder_attention_mask,
							)




@require_flax
class     _UpperCamelCase			(   _A       ,    unittest.TestCase					):






			'''simple docstring'''

			__UpperCamelCase :  Dict    =	True

			__UpperCamelCase :  Optional[int]    =	(
			    (
			        FlaxBertModel,
			        FlaxBertForPreTraining,
			        FlaxBertForMaskedLM,
			        FlaxBertForMultipleChoice,
			        FlaxBertForQuestionAnswering,
			        FlaxBertForNextSentencePrediction,
			        FlaxBertForSequenceClassification,
			        FlaxBertForTokenClassification,
			        FlaxBertForQuestionAnswering,
			    )
			    if is_flax_available()
			    else ()
			)

			def  lowerCAmelCase__		(    self	:    Optional[Any] ):
							UpperCamelCase_:							str		=  FlaxBertModelTester(self )




			@slow
			def  lowerCAmelCase__		(    self	:    str ):
							# Only check this for base model, not necessary for all model classes.
							# This will also help speed-up tests.
							UpperCamelCase_:							Dict		=  FlaxBertModel.from_pretrained("""bert-base-cased""" )
							UpperCamelCase_:							Tuple		=  model(np.ones((1, 1) ) )
							self.assertIsNotNone(snake_case_ )








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

import pyarrow as pa
import pyarrow.json as paj

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


lowerCamelCase_  :	Optional[Any] 				=					datasets.utils.logging.get_logger(__name__)

@dataclass
class     _UpperCamelCase			(   datasets.BuilderConfig					):






			'''simple docstring'''

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




class     _UpperCamelCase			(   datasets.ArrowBasedBuilder					):






			'''simple docstring'''

			__UpperCamelCase :  Tuple    =	JsonConfig

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

			def  lowerCAmelCase__		(    self	:    str  ,		snake_case_	:    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 ):
															UpperCamelCase_:							Union[str, Any]		=  self.config.features.arrow_schema.field(snake_case_ ).type
															UpperCamelCase_:							Tuple		=  pa_table.append_column(snake_case_  ,		pa.array([None] * len(snake_case_ )  ,		type=snake_case_ ) )
											# 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(snake_case_  ,		self.config.features.arrow_schema )
							return pa_table




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

															# We keep only the field we are interested in
															UpperCamelCase_:							Optional[int]		=  dataset[self.config.field]

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

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






