
import inspect
import unittest

import numpy as np

from transformers import BeitConfig
from transformers.testing_utils import require_flax, require_vision, slow
from transformers.utils import cached_property, is_flax_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor


if is_flax_available():
							import jax

							from transformers import FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling, FlaxBeitModel

if is_vision_available():
							from PIL import Image

							from transformers import BeitImageProcessor




class   __lowercase						(						unittest.TestCase       ):






	def __init__(self				: str      ,	snake_case				: List[Any]      ,	snake_case				: int=100      ,	snake_case				: List[str]=13      ,	snake_case				: List[str]=30      ,	snake_case				: Dict=2      ,	snake_case				: int=3      ,	snake_case				: List[str]=True      ,	snake_case				: int=True      ,	snake_case				: Optional[Any]=32      ,	snake_case				: Optional[Any]=5      ,	snake_case				: int=4      ,	snake_case				: Tuple=37      ,	snake_case				: Tuple="gelu"      ,	snake_case				: str=0.1      ,	snake_case				: Optional[int]=0.1      ,	snake_case				: Tuple=10      ,	snake_case				: Any=0.02      ,	snake_case				: Union[str, Any]=3      ,	) ->			Union[str, Any]:
							_lowercase       :	Tuple           =   parent
							_lowercase       :	Union[str, Any]           =   vocab_size
							_lowercase       :	int           =   batch_size
							_lowercase       :	Any           =   image_size
							_lowercase       :	Union[str, Any]           =   patch_size
							_lowercase       :	List[str]           =   num_channels
							_lowercase       :	str           =   is_training
							_lowercase       :	List[Any]           =   use_labels
							_lowercase       :	Tuple           =   hidden_size
							_lowercase       :	Dict           =   num_hidden_layers
							_lowercase       :	Optional[int]           =   num_attention_heads
							_lowercase       :	List[Any]           =   intermediate_size
							_lowercase       :	str           =   hidden_act
							_lowercase       :	Union[str, Any]           =   hidden_dropout_prob
							_lowercase       :	int           =   attention_probs_dropout_prob
							_lowercase       :	Dict           =   type_sequence_label_size
							_lowercase       :	Any           =   initializer_range

							# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
							_lowercase       :	Dict           =   (image_size // patch_size) ** 2
							_lowercase       :	str           =   num_patches + 1






	def       _a(self				: Union[str, Any]    ) ->			int:
							_lowercase       :	Dict           =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

							_lowercase       :	Union[str, Any]           =   None
							if self.use_labels:
													_lowercase       :	Optional[int]           =   ids_tensor([self.batch_size]      ,	self.type_sequence_label_size    )

							_lowercase       :	Tuple           =   BeitConfig(
							    vocab_size=self.vocab_size      ,	image_size=self.image_size      ,	patch_size=self.patch_size      ,	num_channels=self.num_channels      ,	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      ,	is_decoder=snake_case      ,	initializer_range=self.initializer_range      ,	)

							return config, pixel_values, labels






	def       _a(self				: Dict      ,	snake_case				: int      ,	snake_case				: List[str]      ,	snake_case				: List[Any]    ) ->			Union[str, Any]:
							_lowercase       :	Tuple           =   FlaxBeitModel(config=snake_case    )
							_lowercase       :	int           =   model(snake_case    )
							self.parent.assertEqual(result.last_hidden_state.shape      ,	(self.batch_size, self.seq_length, self.hidden_size)    )






	def       _a(self				: Tuple      ,	snake_case				: Optional[int]      ,	snake_case				: Optional[int]      ,	snake_case				: Optional[Any]    ) ->			List[str]:
							_lowercase       :	List[str]           =   FlaxBeitForMaskedImageModeling(config=snake_case    )
							_lowercase       :	Dict           =   model(snake_case    )
							self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.seq_length - 1, self.vocab_size)    )






	def       _a(self				: List[Any]      ,	snake_case				: List[Any]      ,	snake_case				: List[str]      ,	snake_case				: List[str]    ) ->			int:
							_lowercase       :	Any           =   self.type_sequence_label_size
							_lowercase       :	List[Any]           =   FlaxBeitForImageClassification(config=snake_case    )
							_lowercase       :	Optional[int]           =   model(snake_case    )
							self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.type_sequence_label_size)    )

							# test greyscale images
							_lowercase       :	List[Any]           =   1
							_lowercase       :	str           =   FlaxBeitForImageClassification(snake_case    )

							_lowercase       :	Union[str, Any]           =   floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
							_lowercase       :	Optional[int]           =   model(snake_case    )






	def       _a(self				: str    ) ->			List[Any]:
							_lowercase       :	Dict           =   self.prepare_config_and_inputs()
							(
							    (
							    _lowercase
							)      ,    (
							    _lowercase
							)      ,    (
							    _lowercase
							)      ,    
							)       :	Optional[Any]           =   config_and_inputs
							_lowercase       :	List[str]           =   {"pixel_values": pixel_values}
							return config, inputs_dict




@require_flax
class   __lowercase						(						__snake_case   , unittest.TestCase       ):
	_A       =   (
	    (FlaxBeitModel, FlaxBeitForImageClassification, FlaxBeitForMaskedImageModeling) if is_flax_available() else ()
	)






	def       _a(self				: int    ) ->			None:
							_lowercase       :	List[str]           =   FlaxBeitModelTester(self    )
							_lowercase       :	Dict           =   ConfigTester(self      ,	config_class=snake_case      ,	has_text_modality=snake_case      ,	hidden_size=37    )






	def       _a(self				: Any    ) ->			Union[str, Any]:
							self.config_tester.run_common_tests()






	def       _a(self				: Optional[Any]    ) ->			List[str]:
							_lowercase      ,    _lowercase       :	Optional[Any]           =   self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													_lowercase       :	Dict           =   model_class(snake_case    )
													_lowercase       :	Any           =   inspect.signature(model.__call__    )
													# signature.parameters is an OrderedDict => so arg_names order is deterministic
													_lowercase       :	Union[str, Any]           =   [*signature.parameters.keys()]

													_lowercase       :	Dict           =   ["pixel_values"]
													self.assertListEqual(arg_names[:1]      ,	snake_case    )






	def       _a(self				: str    ) ->			Tuple:
							_lowercase      ,    _lowercase       :	int           =   self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													with self.subTest(model_class.__name__    ):
																			_lowercase       :	Optional[int]           =   self._prepare_for_class(snake_case      ,	snake_case    )
																			_lowercase       :	Union[str, Any]           =   model_class(snake_case    )

																			@jax.jit
																			def model_jitted(snake_case				: Any      ,	**snake_case				: Optional[int]    ):
																									return model(pixel_values=snake_case      ,	**snake_case    )

																			with self.subTest("JIT Enabled"    ):
																									_lowercase       :	Tuple           =   model_jitted(**snake_case    ).to_tuple()

																			with self.subTest("JIT Disabled"    ):
																									with jax.disable_jit():
																															_lowercase       :	Dict           =   model_jitted(**snake_case    ).to_tuple()

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






	def       _a(self				: Tuple    ) ->			Union[str, Any]:
							_lowercase       :	Optional[Any]           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*snake_case    )






	def       _a(self				: Union[str, Any]    ) ->			int:
							_lowercase       :	Tuple           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*snake_case    )






	def       _a(self				: List[str]    ) ->			int:
							_lowercase       :	List[str]           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*snake_case    )






	@slow
	def       _a(self				: int    ) ->			int:
							for model_class_name in self.all_model_classes:
													_lowercase       :	Optional[Any]           =   model_class_name.from_pretrained("microsoft/beit-base-patch16-224"    )
													_lowercase       :	Tuple           =   model(np.ones((1, 3, 224, 224)    )    )
													self.assertIsNotNone(snake_case    )



def 		UpperCamelCase     (       )   ->			Any:






						'''simple docstring'''

						_lowercase       :	Optional[int]           =   Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"  )
						return image




@require_vision
@require_flax
class   __lowercase						(						unittest.TestCase       ):






	@cached_property
	def       _a(self				: List[Any]    ) ->			Optional[Any]:
							return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224"    ) if is_vision_available() else None






	@slow
	def       _a(self				: Optional[Any]    ) ->			List[Any]:
							_lowercase       :	str           =   FlaxBeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"    )

							_lowercase       :	Optional[int]           =   self.default_image_processor
							_lowercase       :	List[str]           =   prepare_img()
							_lowercase       :	Tuple           =   image_processor(images=snake_case      ,	return_tensors="np"    ).pixel_values

							# prepare bool_masked_pos
							_lowercase       :	Tuple           =   np.ones((1, 196)      ,	dtype=snake_case    )

							# forward pass
							_lowercase       :	Optional[int]           =   model(pixel_values=snake_case      ,	bool_masked_pos=snake_case    )
							_lowercase       :	int           =   outputs.logits

							# verify the logits
							_lowercase       :	Optional[int]           =   (1, 196, 8192)
							self.assertEqual(logits.shape      ,	snake_case    )

							_lowercase       :	Union[str, Any]           =   np.array(
							    [[-3.24_37, 0.50_72, -13.91_74], [-3.24_56, 0.49_48, -13.94_01], [-3.20_33, 0.51_21, -13.85_50]]    )

							self.assertTrue(np.allclose(logits[bool_masked_pos][:3, :3]      ,	snake_case      ,	atol=1e-2    )    )






	@slow
	def       _a(self				: List[str]    ) ->			Any:
							_lowercase       :	Dict           =   FlaxBeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"    )

							_lowercase       :	Optional[int]           =   self.default_image_processor
							_lowercase       :	Dict           =   prepare_img()
							_lowercase       :	List[Any]           =   image_processor(images=snake_case      ,	return_tensors="np"    )

							# forward pass
							_lowercase       :	int           =   model(**snake_case    )
							_lowercase       :	int           =   outputs.logits

							# verify the logits
							_lowercase       :	List[str]           =   (1, 1000)
							self.assertEqual(logits.shape      ,	snake_case    )

							_lowercase       :	Union[str, Any]           =   np.array([-1.23_85, -1.09_87, -1.01_08]    )

							self.assertTrue(np.allclose(logits[0, :3]      ,	snake_case      ,	atol=1e-4    )    )

							_lowercase       :	Dict           =   281
							self.assertEqual(logits.argmax(-1    ).item()      ,	snake_case    )






	@slow
	def       _a(self				: Optional[Any]    ) ->			Tuple:
							_lowercase       :	List[Any]           =   FlaxBeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"    )

							_lowercase       :	Dict           =   self.default_image_processor
							_lowercase       :	int           =   prepare_img()
							_lowercase       :	Tuple           =   image_processor(images=snake_case      ,	return_tensors="np"    )

							# forward pass
							_lowercase       :	Union[str, Any]           =   model(**snake_case    )
							_lowercase       :	List[Any]           =   outputs.logits

							# verify the logits
							_lowercase       :	List[str]           =   (1, 2_1841)
							self.assertEqual(logits.shape      ,	snake_case    )

							_lowercase       :	Optional[Any]           =   np.array([1.68_81, -0.27_87, 0.59_01]    )

							self.assertTrue(np.allclose(logits[0, :3]      ,	snake_case      ,	atol=1e-4    )    )

							_lowercase       :	Optional[Any]           =   2396
							self.assertEqual(logits.argmax(-1    ).item()      ,	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


UpperCamelCase_						:     Union[str, Any]	       =							datasets.utils.logging.get_logger(__name__)




@dataclass
class   __lowercase						(						datasets.BuilderConfig       ):
	_A       =   None
	_A       =   "utf-8"
	_A       =   None
	_A       =   None
	_A       =   True  # deprecated
	_A       =   None  # deprecated
	_A       =   10 << 20  # 10MB
	_A       =   None




class   __lowercase						(						datasets.ArrowBasedBuilder       ):
	_A       =   JsonConfig






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






	def       _a(self				: List[Any]      ,	snake_case				: 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       :	Any           =   self.config.features.arrow_schema.field(snake_case    ).type
																			_lowercase       :	Optional[Any]           =   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
													_lowercase       :	str           =   table_cast(snake_case      ,	self.config.features.arrow_schema    )
							return pa_table






	def       _a(self				: str      ,	snake_case				: Tuple    ) ->			Tuple:
							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:
																									_lowercase       :	Any           =   json.load(snake_case    )

																			# We keep only the field we are interested in
																			_lowercase       :	List[Any]           =   dataset[self.config.field]

																			# We accept two format: a list of dicts or a dict of lists
																			if isinstance(snake_case      ,	(list, tuple)    ):
																									_lowercase       :	Optional[int]           =   set().union(*[row.keys() for row in dataset]    )
																									_lowercase       :	Union[str, Any]           =   {col: [row.get(snake_case    ) for row in dataset] for col in keys}
																			else:
																									_lowercase       :	List[Any]           =   dataset
																			_lowercase       :	int           =   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:
																									_lowercase       :	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
																									_lowercase       :	Dict           =   max(self.config.chunksize // 32      ,	16 << 10    )
																									_lowercase       :	int           =   (
																									    self.config.encoding_errors if self.config.encoding_errors is not None else "strict"
																									)
																									while True:
																															_lowercase       :	Tuple           =   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":
																																					_lowercase       :	Union[str, Any]           =   batch.decode(self.config.encoding      ,	errors=snake_case    ).encode("utf-8"    )
																															try:
																																					while True:
																																											try:
																																																	_lowercase       :	Optional[Any]           =   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:
																																																	_lowercase       :	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:
																																																	_lowercase       :	int           =   set().union(*[row.keys() for row in dataset]    )
																																																	_lowercase       :	Dict           =   {col: [row.get(snake_case    ) for row in dataset] for col in keys}
																																																	_lowercase       :	str           =   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




