



import inspect
import unittest

from transformers import ViTHybridConfig
from transformers.testing_utils import require_accelerate, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch
							from torch import nn

							from transformers import ViTHybridForImageClassification, ViTHybridImageProcessor, ViTHybridModel
							from transformers.models.vit_hybrid.modeling_vit_hybrid import VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							from PIL import Image


class   __lowerCAmelCase  :
					"""simple docstring"""






					def __init__(  self       :							List[Any]					,       _snake_case       :							Dict					,       _snake_case       :							Union[str, Any]=13					,       _snake_case       :							Tuple=64					,       _snake_case       :							Any=2					,       _snake_case       :							Optional[Any]=3					,       _snake_case       :							Any=True					,       _snake_case       :							Dict=True					,       _snake_case       :							Optional[int]=32					,       _snake_case       :							List[str]=5					,       _snake_case       :							Union[str, Any]=4					,       _snake_case       :							str=37					,       _snake_case       :							str="gelu"					,       _snake_case       :							Dict=0.1					,       _snake_case       :							Union[str, Any]=0.1					,       _snake_case       :							Union[str, Any]=10					,       _snake_case       :							Optional[int]=0.02					,       _snake_case       :							Optional[Any]=[1, 16, 4, 4]					,       _snake_case       :							Optional[Any]=None					,       ):






								"""simple docstring"""


								A__	  =     parent
								A__	  =     batch_size
								A__	  =     image_size
								A__	  =     patch_size
								A__	  =     num_channels
								A__	  =     is_training
								A__	  =     use_labels
								A__	  =     hidden_size
								A__	  =     num_hidden_layers
								A__	  =     num_attention_heads
								A__	  =     intermediate_size
								A__	  =     hidden_act
								A__	  =     hidden_dropout_prob
								A__	  =     attention_probs_dropout_prob
								A__	  =     type_sequence_label_size
								A__	  =     initializer_range
								A__	  =     scope
								A__	  =     backbone_featmap_shape

								# in ViT hybrid, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
								# the number of patches is based on the feature map of the backbone, which by default uses an output stride
								# of 32, which means that the feature map has a spatial resolution of 1/32 of the input image size
								A__	  =     (self.image_size // 32) ** 2
								A__	  =     num_patches + 1



					def      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


								A__	  =     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

								A__	  =     None
								if self.use_labels:
											A__	  =     ids_tensor([self.batch_size]					,       self.type_sequence_label_size       )

								A__	  =     self.get_config()

								return config, pixel_values, labels



					def      _a				(  self       :							Any       ):






								"""simple docstring"""


								A__	  =     {
								    'global_padding': 'same',
								    'layer_type': 'bottleneck',
								    'depths': [3, 4, 9],
								    'out_features': ['stage1', 'stage2', 'stage3'],
								    'embedding_dynamic_padding': True,
								    'hidden_sizes': [4, 8, 16, 32],
								    'num_groups': 2,
								}

								return ViTHybridConfig(
								    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					,       backbone_featmap_shape=self.backbone_featmap_shape					,       backbone_config=_snake_case					,       )



					def      _a				(  self       :							Union[str, Any]					,       _snake_case       :							int					,       _snake_case       :							Optional[int]					,       _snake_case       :							int       ):






								"""simple docstring"""


								A__	  =     ViTHybridModel(config=_snake_case       )
								model.to(_snake_case       )
								model.eval()
								A__	  =     model(_snake_case       )
								self.parent.assertEqual(result.last_hidden_state.shape					,       (self.batch_size, self.seq_length, self.hidden_size)       )



					def      _a				(  self       :							int					,       _snake_case       :							Dict					,       _snake_case       :							Dict					,       _snake_case       :							Optional[Any]       ):






								"""simple docstring"""


								A__	  =     self.type_sequence_label_size
								A__	  =     ViTHybridForImageClassification(_snake_case       )
								model.to(_snake_case       )
								model.eval()
								A__	  =     model(_snake_case					,       labels=_snake_case       )
								self.parent.assertEqual(result.logits.shape					,       (self.batch_size, self.type_sequence_label_size)       )



					def      _a				(  self       :							Any       ):






								"""simple docstring"""


								A__	  =     self.prepare_config_and_inputs()
								A__  ,			A__  ,			A__	  =     config_and_inputs
								A__	  =     {'pixel_values': pixel_values}
								return config, inputs_dict







@require_torch
class   __lowerCAmelCase  (		UpperCAmelCase_				,  UpperCAmelCase_				,  unittest.TestCase   ):
					"""simple docstring"""






					A__					:		Any      							=     (ViTHybridModel, ViTHybridForImageClassification) if is_torch_available() else ()
					A__					:		str      							=     (
					    {"feature-extraction": ViTHybridModel, "image-classification": ViTHybridForImageClassification}
					    if is_torch_available()
					    else {}
					)
					A__					:		int      							=     False
					A__					:		Dict      							=     False
					A__					:		Tuple      							=     False



					def      _a				(  self       :							List[str]       ):






								"""simple docstring"""


								A__	  =     ViTHybridModelTester(self       )
								A__	  =     ConfigTester(self					,       config_class=_snake_case					,       has_text_modality=_snake_case					,       hidden_size=37       )



					def      _a				(  self       :							int       ):






								"""simple docstring"""


								self.config_tester.run_common_tests()



					@unittest.skip(reason='ViT does not use inputs_embeds'       )
					def      _a				(  self       :							Optional[int]       ):






								"""simple docstring"""


								pass



					def      _a				(  self       :							int       ):






								"""simple docstring"""


								A__  ,			A__	  =     self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											A__	  =     model_class(_snake_case       )
											self.assertIsInstance(model.get_input_embeddings()					,       (nn.Module)       )
											A__	  =     model.get_output_embeddings()
											self.assertTrue(x is None or isinstance(_snake_case					,       nn.Linear       )       )



					def      _a				(  self       :							Dict       ):






								"""simple docstring"""


								A__  ,			A__	  =     self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											A__	  =     model_class(_snake_case       )
											A__	  =     inspect.signature(model.forward       )
											# signature.parameters is an OrderedDict => so arg_names order is deterministic
											A__	  =     [*signature.parameters.keys()]

											A__	  =     ['pixel_values']
											self.assertListEqual(arg_names[:1]					,       _snake_case       )



					def      _a				(  self       :							List[str]       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_snake_case       )



					def      _a				(  self       :							List[Any]       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*_snake_case       )



					def      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


								A__  ,			A__	  =     self.model_tester.prepare_config_and_inputs_for_common()

								A__	  =     _config_zero_init(_snake_case       )
								for model_class in self.all_model_classes:
											A__	  =     model_class(config=_snake_case       )
											# Skip the check for the backbone
											for name, module in model.named_modules():
														if module.__class__.__name__ == "ViTHybridPatchEmbeddings":
																	A__	  =     [F'''{name}.{key}''' for key in module.state_dict().keys()]
																	break

											for name, param in model.named_parameters():
														if param.requires_grad:
																	if name in backbone_params:
																				continue
																	self.assertIn(
																	    ((param.data.mean() * 1E9).round() / 1E9).item()					,       [0.0, 1.0]					,       msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''					,       )



					@slow
					def      _a				(  self       :							Any       ):






								"""simple docstring"""


								for model_name in VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											A__	  =     ViTHybridModel.from_pretrained(_snake_case       )
											self.assertIsNotNone(_snake_case       )







def A    (							)    ->   Optional[Any]:
			A__	  =     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'       )
			return image






@require_torch
@require_vision
class   __lowerCAmelCase  (		unittest.TestCase   ):
					"""simple docstring"""






					@cached_property
					def      _a				(  self       :							int       ):






								"""simple docstring"""


								return (
								    ViTHybridImageProcessor.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]       )
								    if is_vision_available()
								    else None
								)



					@slow
					def      _a				(  self       :							Any       ):






								"""simple docstring"""


								A__	  =     ViTHybridForImageClassification.from_pretrained(VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST[0]       ).to(
								    _snake_case       )

								A__	  =     self.default_image_processor
								A__	  =     prepare_img()
								A__	  =     image_processor(images=_snake_case					,       return_tensors='pt'       ).to(_snake_case       )

								# forward pass
								with torch.no_grad():
											A__	  =     model(**_snake_case       )

								# verify the logits
								A__	  =     torch.Size((1, 10_00)       )
								self.assertEqual(outputs.logits.shape					,       _snake_case       )

								A__	  =     torch.tensor([-1.9090, -0.4993, -0.2389]       ).to(_snake_case       )

								self.assertTrue(torch.allclose(outputs.logits[0, :3]					,       _snake_case					,       atol=1E-4       )       )



					@slow
					@require_accelerate
					def      _a				(  self       :							str       ):






								"""simple docstring"""


								A__	  =     ViTHybridImageProcessor.from_pretrained('google/vit-hybrid-base-bit-384'       )
								A__	  =     ViTHybridForImageClassification.from_pretrained('google/vit-hybrid-base-bit-384'					,       device_map='auto'       )

								A__	  =     prepare_img()

								A__	  =     image_processor(images=_snake_case					,       return_tensors='pt'       )
								A__	  =     model(**_snake_case       )
								A__	  =     outputs.logits
								# model predicts one of the 1000 ImageNet classes
								A__	  =     logits.argmax(-1       ).item()

								self.assertTrue(model.config.idalabel[predicted_class_idx]					,       'tabby, tabby cat'       )






"""simple docstring"""



import copy
import os
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import pytest

from datasets.arrow_writer import ArrowWriter, OptimizedTypedSequence, ParquetWriter, TypedSequence
from datasets.features import ArrayaD, ClassLabel, Features, Image, Value
from datasets.features.features import ArrayaDExtensionType, cast_to_python_objects
from datasets.keyhash import DuplicatedKeysError, InvalidKeyError

from .utils import require_pil




class 		a_    (			_UpperCAmelCase						):







    def   _snake_case       (      self      :							int  )		->Union[str, Any]:






        '''simple docstring'''






        _UpperCAmelCase        = pa.array(TypedSequence([1, 2, 3]  )  )
        self.assertEqual(arr.type  ,       pa.intaa()  )







    def   _snake_case       (      self      :							List[Any]  )		->Any:






        '''simple docstring'''






        with self.assertRaises(__UpperCamelCase  ):
            _UpperCAmelCase        = pa.array(TypedSequence([1, 2, 3]  )  ,       type=pa.intaa()  )







    def   _snake_case       (      self      :							Tuple  )		->int:






        '''simple docstring'''






        with self.assertRaises(__UpperCamelCase  ):
            _UpperCAmelCase        = pa.array(TypedSequence([1, 2, 3]  ,       try_type=Value("""bool"""  )  ,       type=Value("""int64"""  )  )  )







    def   _snake_case       (      self      :							Optional[Any]  )		->Tuple:






        '''simple docstring'''






        _UpperCAmelCase        = pa.array(TypedSequence([1, 2, 3]  ,       type=Value("""int32"""  )  )  )
        self.assertEqual(arr.type  ,       pa.intaa()  )







    def   _snake_case       (      self      :							List[Any]  )		->int:






        '''simple docstring'''






        with self.assertRaises((TypeError, pa.lib.ArrowInvalid)  ):
            _UpperCAmelCase        = pa.array(TypedSequence(["""foo""", """bar"""]  ,       type=Value("""int64"""  )  )  )







    def   _snake_case       (      self      :							int  )		->Tuple:






        '''simple docstring'''






        _UpperCAmelCase        = pa.array(TypedSequence([1, 2, 3]  ,       try_type=Value("""int32"""  )  )  )
        self.assertEqual(arr.type  ,       pa.intaa()  )







    def   _snake_case       (      self      :							Dict  )		->str:






        '''simple docstring'''






        _UpperCAmelCase        = pa.array(TypedSequence(["""foo""", """bar"""]  ,       try_type=Value("""int64"""  )  )  )
        self.assertEqual(arr.type  ,       pa.string()  )







    def   _snake_case       (      self      :							Union[str, Any]  )		->Optional[int]:






        '''simple docstring'''






        _UpperCAmelCase        = pa.array(TypedSequence([[[1, 2, 3]]]  ,       type=ArrayaD((1, 3)  ,       """int64"""  )  )  )
        self.assertEqual(arr.type  ,       ArrayaDExtensionType((1, 3)  ,       """int64"""  )  )







    def   _snake_case       (      self      :							Tuple  )		->Any:






        '''simple docstring'''






        with self.assertRaises((TypeError, pa.lib.ArrowInvalid)  ):
            _UpperCAmelCase        = pa.array(TypedSequence(["""foo""", """bar"""]  ,       type=ArrayaD((1, 3)  ,       """int64"""  )  )  )







    def   _snake_case       (      self      :							List[Any]  )		->List[Any]:






        '''simple docstring'''






        _UpperCAmelCase        = pa.array(TypedSequence([[[1, 2, 3]]]  ,       try_type=ArrayaD((1, 3)  ,       """int64"""  )  )  )
        self.assertEqual(arr.type  ,       ArrayaDExtensionType((1, 3)  ,       """int64"""  )  )







    def   _snake_case       (      self      :							Any  )		->Dict:






        '''simple docstring'''






        _UpperCAmelCase        = pa.array(TypedSequence(["""foo""", """bar"""]  ,       try_type=ArrayaD((1, 3)  ,       """int64"""  )  )  )
        self.assertEqual(arr.type  ,       pa.string()  )







    @require_pil
    def   _snake_case       (      self      :							Optional[Any]  )		->Tuple:






        '''simple docstring'''






        import PIL.Image

        _UpperCAmelCase        = PIL.Image.fromarray(np.arange(10  ,       dtype=np.uinta  ).reshape(2  ,       5  )  )
        with patch(
            """datasets.arrow_writer.cast_to_python_objects"""  ,       side_effect=__UpperCamelCase  ) as mock_cast_to_python_objects:
            _UpperCAmelCase        = pa.array(TypedSequence([{"""path""": None, """bytes""": b"""image_bytes"""}, pil_image]  ,       type=Image()  )  )
            _UpperCAmelCase     ,_UpperCAmelCase        = mock_cast_to_python_objects.call_args_list[-1]
            self.assertIn("""optimize_list_casting"""  ,       __UpperCamelCase  )
            self.assertFalse(kwargs["""optimize_list_casting"""]  )







def 			_UpperCamelCase						(    _A       ,   _A   )     ->	Dict:





    """simple docstring"""

    _UpperCAmelCase        = pa.BufferReader(_A   ) if isinstance(_A       ,   pa.Buffer   ) else pa.memory_map(_A   )
    _UpperCAmelCase        = pa.ipc.open_stream(_A   )
    _UpperCAmelCase        = f.read_all()
    assert len(pa_table.to_batches()   ) == expected_num_chunks
    assert pa_table.to_pydict() == {"col_1": ["foo", "bar"], "col_2": [1, 2]}
    del pa_table


@pytest.mark.parametrize("""writer_batch_size"""       ,   [None, 1, 1_0]   )
@pytest.mark.parametrize(
    """fields"""       ,   [None, {"""col_1""": pa.string(), """col_2""": pa.intaa()}, {"""col_1""": pa.string(), """col_2""": pa.intaa()}]   )
def 			_UpperCamelCase						(    _A       ,   _A   )     ->	Any:





    """simple docstring"""

    _UpperCAmelCase        = pa.BufferOutputStream()
    _UpperCAmelCase        = pa.schema(_A   ) if fields else None
    with ArrowWriter(stream=_A       ,   schema=_A       ,   writer_batch_size=_A   ) as writer:
        writer.write({"""col_1""": """foo""", """col_2""": 1}   )
        writer.write({"""col_1""": """bar""", """col_2""": 2}   )
        _UpperCAmelCase     ,_UpperCAmelCase        = writer.finalize()
    assert num_examples == 2
    assert num_bytes > 0
    if not fields:
        _UpperCAmelCase        = {"""col_1""": pa.string(), """col_2""": pa.intaa()}
    assert writer._schema == pa.schema(_A       ,   metadata=writer._schema.metadata   )
    _check_output(output.getvalue()       ,   expected_num_chunks=num_examples if writer_batch_size == 1 else 1   )


def 			_UpperCamelCase						(    )     ->	Dict:





    """simple docstring"""

    _UpperCAmelCase        = pa.BufferOutputStream()
    _UpperCAmelCase        = Features({"""labels""": ClassLabel(names=["""neg""", """pos"""]   )}   )
    with ArrowWriter(stream=_A       ,   features=_A   ) as writer:
        writer.write({"""labels""": 0}   )
        writer.write({"""labels""": 1}   )
        _UpperCAmelCase     ,_UpperCAmelCase        = writer.finalize()
    assert num_examples == 2
    assert num_bytes > 0
    assert writer._schema == features.arrow_schema
    assert writer._schema.metadata == features.arrow_schema.metadata
    _UpperCAmelCase        = pa.BufferReader(output.getvalue()   )
    _UpperCAmelCase        = pa.ipc.open_stream(_A   )
    _UpperCAmelCase        = f.read_all()
    _UpperCAmelCase        = pa_table.schema
    assert pa_table.num_rows == 2
    assert schema == features.arrow_schema
    assert schema.metadata == features.arrow_schema.metadata
    assert features == Features.from_arrow_schema(_A   )


@pytest.mark.parametrize("""writer_batch_size"""       ,   [None, 1, 1_0]   )
def 			_UpperCamelCase						(    _A   )     ->	int:





    """simple docstring"""

    _UpperCAmelCase        = pa.BufferOutputStream()
    with ArrowWriter(
        stream=_A       ,   writer_batch_size=_A       ,   hash_salt="""split_name"""       ,   check_duplicates=_A       ,   ) as writer:
        with pytest.raises(_A   ):
            writer.write({"""col_1""": """foo""", """col_2""": 1}       ,   key=[1, 2]   )
            _UpperCAmelCase     ,_UpperCAmelCase        = writer.finalize()


@pytest.mark.parametrize("""writer_batch_size"""       ,   [None, 2, 1_0]   )
def 			_UpperCamelCase						(    _A   )     ->	List[Any]:





    """simple docstring"""

    _UpperCAmelCase        = pa.BufferOutputStream()
    with ArrowWriter(
        stream=_A       ,   writer_batch_size=_A       ,   hash_salt="""split_name"""       ,   check_duplicates=_A       ,   ) as writer:
        with pytest.raises(_A   ):
            writer.write({"""col_1""": """foo""", """col_2""": 1}       ,   key=1_0   )
            writer.write({"""col_1""": """bar""", """col_2""": 2}       ,   key=1_0   )
            _UpperCAmelCase     ,_UpperCAmelCase        = writer.finalize()


@pytest.mark.parametrize("""writer_batch_size"""       ,   [None, 2, 1_0]   )
def 			_UpperCamelCase						(    _A   )     ->	Dict:





    """simple docstring"""

    _UpperCAmelCase        = pa.BufferOutputStream()
    with ArrowWriter(
        stream=_A       ,   writer_batch_size=_A       ,   hash_salt="""split_name"""       ,   check_duplicates=_A       ,   ) as writer:
        writer.write({"""col_1""": """foo""", """col_2""": 1}       ,   key=1   )
        writer.write({"""col_1""": """bar""", """col_2""": 2}       ,   key=2   )
        _UpperCAmelCase     ,_UpperCAmelCase        = writer.finalize()
    assert num_examples == 2
    assert num_bytes > 0
    _check_output(output.getvalue()       ,   expected_num_chunks=num_examples if writer_batch_size == 1 else 1   )


@pytest.mark.parametrize("""writer_batch_size"""       ,   [None, 1, 1_0]   )
@pytest.mark.parametrize(
    """fields"""       ,   [None, {"""col_1""": pa.string(), """col_2""": pa.intaa()}, {"""col_1""": pa.string(), """col_2""": pa.intaa()}]   )
def 			_UpperCamelCase						(    _A       ,   _A   )     ->	Any:





    """simple docstring"""

    _UpperCAmelCase        = pa.BufferOutputStream()
    _UpperCAmelCase        = pa.schema(_A   ) if fields else None
    with ArrowWriter(stream=_A       ,   schema=_A       ,   writer_batch_size=_A   ) as writer:
        writer.write_batch({"""col_1""": ["""foo""", """bar"""], """col_2""": [1, 2]}   )
        writer.write_batch({"""col_1""": [], """col_2""": []}   )
        _UpperCAmelCase     ,_UpperCAmelCase        = writer.finalize()
    assert num_examples == 2
    assert num_bytes > 0
    if not fields:
        _UpperCAmelCase        = {"""col_1""": pa.string(), """col_2""": pa.intaa()}
    assert writer._schema == pa.schema(_A       ,   metadata=writer._schema.metadata   )
    _check_output(output.getvalue()       ,   expected_num_chunks=num_examples if writer_batch_size == 1 else 1   )


@pytest.mark.parametrize("""writer_batch_size"""       ,   [None, 1, 1_0]   )
@pytest.mark.parametrize(
    """fields"""       ,   [None, {"""col_1""": pa.string(), """col_2""": pa.intaa()}, {"""col_1""": pa.string(), """col_2""": pa.intaa()}]   )
def 			_UpperCamelCase						(    _A       ,   _A   )     ->	str:





    """simple docstring"""

    _UpperCAmelCase        = pa.BufferOutputStream()
    _UpperCAmelCase        = pa.schema(_A   ) if fields else None
    with ArrowWriter(stream=_A       ,   schema=_A       ,   writer_batch_size=_A   ) as writer:
        writer.write_table(pa.Table.from_pydict({"""col_1""": ["""foo""", """bar"""], """col_2""": [1, 2]}   )   )
        _UpperCAmelCase     ,_UpperCAmelCase        = writer.finalize()
    assert num_examples == 2
    assert num_bytes > 0
    if not fields:
        _UpperCAmelCase        = {"""col_1""": pa.string(), """col_2""": pa.intaa()}
    assert writer._schema == pa.schema(_A       ,   metadata=writer._schema.metadata   )
    _check_output(output.getvalue()       ,   expected_num_chunks=num_examples if writer_batch_size == 1 else 1   )


@pytest.mark.parametrize("""writer_batch_size"""       ,   [None, 1, 1_0]   )
@pytest.mark.parametrize(
    """fields"""       ,   [None, {"""col_1""": pa.string(), """col_2""": pa.intaa()}, {"""col_1""": pa.string(), """col_2""": pa.intaa()}]   )
def 			_UpperCamelCase						(    _A       ,   _A   )     ->	int:





    """simple docstring"""

    _UpperCAmelCase        = pa.BufferOutputStream()
    _UpperCAmelCase        = pa.schema(_A   ) if fields else None
    with ArrowWriter(stream=_A       ,   schema=_A       ,   writer_batch_size=_A   ) as writer:
        writer.write_row(pa.Table.from_pydict({"""col_1""": ["""foo"""], """col_2""": [1]}   )   )
        writer.write_row(pa.Table.from_pydict({"""col_1""": ["""bar"""], """col_2""": [2]}   )   )
        _UpperCAmelCase     ,_UpperCAmelCase        = writer.finalize()
    assert num_examples == 2
    assert num_bytes > 0
    if not fields:
        _UpperCAmelCase        = {"""col_1""": pa.string(), """col_2""": pa.intaa()}
    assert writer._schema == pa.schema(_A       ,   metadata=writer._schema.metadata   )
    _check_output(output.getvalue()       ,   expected_num_chunks=num_examples if writer_batch_size == 1 else 1   )


def 			_UpperCamelCase						(    )     ->	Dict:





    """simple docstring"""

    with tempfile.TemporaryDirectory() as tmp_dir:
        _UpperCAmelCase        = {"""col_1""": pa.string(), """col_2""": pa.intaa()}
        _UpperCAmelCase        = os.path.join(_A       ,   """test.arrow"""   )
        with ArrowWriter(path=_A       ,   schema=pa.schema(_A   )   ) as writer:
            writer.write_batch({"""col_1""": ["""foo""", """bar"""], """col_2""": [1, 2]}   )
            _UpperCAmelCase     ,_UpperCAmelCase        = writer.finalize()
        assert num_examples == 2
        assert num_bytes > 0
        assert writer._schema == pa.schema(_A       ,   metadata=writer._schema.metadata   )
        _check_output(_A       ,   1   )


def 			_UpperCamelCase						(    _A   )     ->	Optional[int]:





    """simple docstring"""

    if pa.types.is_list(_A   ):
        return get_base_dtype(arr_type.value_type   )
    else:
        return arr_type


def 			_UpperCamelCase						(    _A       ,   _A   )     ->	str:





    """simple docstring"""

    if isinstance(lst[0]       ,   _A   ):
        change_first_primitive_element_in_list(lst[0]       ,   _A   )
    else:
        _UpperCAmelCase        = value


@pytest.mark.parametrize("""optimized_int_type, expected_dtype"""       ,   [(None, pa.intaa()), (Value("""int32"""   ), pa.intaa())]   )
@pytest.mark.parametrize("""sequence"""       ,   [[1, 2, 3], [[1, 2, 3]], [[[1, 2, 3]]]]   )
def 			_UpperCamelCase						(    _A       ,   _A       ,   _A   )     ->	Optional[int]:





    """simple docstring"""

    _UpperCAmelCase        = pa.array(TypedSequence(_A       ,   optimized_int_type=_A   )   )
    assert get_base_dtype(arr.type   ) == expected_dtype


@pytest.mark.parametrize(
    """col, expected_dtype"""       ,   [
        ("""attention_mask""", pa.inta()),
        ("""special_tokens_mask""", pa.inta()),
        ("""token_type_ids""", pa.inta()),
        ("""input_ids""", pa.intaa()),
        ("""other""", pa.intaa()),
    ]       ,   )
@pytest.mark.parametrize("""sequence"""       ,   [[1, 2, 3], [[1, 2, 3]], [[[1, 2, 3]]]]   )
def 			_UpperCamelCase						(    _A       ,   _A       ,   _A   )     ->	Any:





    """simple docstring"""

    _UpperCAmelCase        = pa.array(OptimizedTypedSequence(_A       ,   col=_A   )   )
    assert get_base_dtype(arr.type   ) == expected_dtype

    # not in range
    if col != "other":
        # avoids errors due to in-place modifications
        _UpperCAmelCase        = copy.deepcopy(_A   )
        _UpperCAmelCase        = np.iinfo(expected_dtype.to_pandas_dtype()   ).max + 1
        change_first_primitive_element_in_list(_A       ,   _A   )
        _UpperCAmelCase        = pa.array(OptimizedTypedSequence(_A       ,   col=_A   )   )
        assert get_base_dtype(arr.type   ) == pa.intaa()


@pytest.mark.parametrize("""raise_exception"""       ,   [False, True]   )
def 			_UpperCamelCase						(    _A       ,   _A   )     ->	Optional[Any]:





    """simple docstring"""

    _UpperCAmelCase        = str(tmp_path / """dataset-train.arrow"""   )
    try:
        with ArrowWriter(path=_A   ) as writer:
            if raise_exception:
                raise pa.lib.ArrowInvalid()
            else:
                writer.stream.close()
    except pa.lib.ArrowInvalid:
        pass
    finally:
        assert writer.stream.closed


def 			_UpperCamelCase						(    _A   )     ->	List[str]:





    """simple docstring"""

    _UpperCAmelCase        = """mock://dataset-train.arrow"""
    with ArrowWriter(path=_A       ,   storage_options=mockfs.storage_options   ) as writer:
        assert isinstance(writer._fs       ,   type(_A   )   )
        assert writer._fs.storage_options == mockfs.storage_options
        writer.write({"""col_1""": """foo""", """col_2""": 1}   )
        writer.write({"""col_1""": """bar""", """col_2""": 2}   )
        _UpperCAmelCase     ,_UpperCAmelCase        = writer.finalize()
    assert num_examples == 2
    assert num_bytes > 0
    assert mockfs.exists(_A   )


def 			_UpperCamelCase						(    )     ->	List[Any]:





    """simple docstring"""

    _UpperCAmelCase        = pa.BufferOutputStream()
    with ParquetWriter(stream=_A   ) as writer:
        writer.write({"""col_1""": """foo""", """col_2""": 1}   )
        writer.write({"""col_1""": """bar""", """col_2""": 2}   )
        _UpperCAmelCase     ,_UpperCAmelCase        = writer.finalize()
    assert num_examples == 2
    assert num_bytes > 0
    _UpperCAmelCase        = pa.BufferReader(output.getvalue()   )
    _UpperCAmelCase        = pq.read_table(_A   )
    assert pa_table.to_pydict() == {"col_1": ["foo", "bar"], "col_2": [1, 2]}


@require_pil
@pytest.mark.parametrize("""embed_local_files"""       ,   [False, True]   )
def 			_UpperCamelCase						(    _A       ,   _A   )     ->	Union[str, Any]:





    """simple docstring"""

    import PIL.Image

    _UpperCAmelCase        = str(tmp_path / """test_image_rgb.jpg"""   )
    PIL.Image.fromarray(np.zeros((5, 5)       ,   dtype=np.uinta   )   ).save(_A       ,   format="""png"""   )
    _UpperCAmelCase        = pa.BufferOutputStream()
    with ParquetWriter(
        stream=_A       ,   features=Features({"""image""": Image()}   )       ,   embed_local_files=_A   ) as writer:
        writer.write({"""image""": image_path}   )
        writer.finalize()
    _UpperCAmelCase        = pa.BufferReader(output.getvalue()   )
    _UpperCAmelCase        = pq.read_table(_A   )
    _UpperCAmelCase        = pa_table.to_pydict()
    if embed_local_files:
        assert isinstance(out["""image"""][0]["""path"""]       ,   _A   )
        with open(_A       ,   """rb"""   ) as f:
            assert out["image"][0]["bytes"] == f.read()
    else:
        assert out["image"][0]["path"] == image_path
        assert out["image"][0]["bytes"] is None







def 			_UpperCamelCase						(    )     ->	Tuple:





    """simple docstring"""

    _UpperCAmelCase        = pa.schema([pa.field("""col_1"""       ,   pa.string()       ,   nullable=_A   )]   )
    _UpperCAmelCase        = pa.BufferOutputStream()
    with ArrowWriter(stream=_A   ) as writer:
        writer._build_writer(inferred_schema=_A   )
        assert writer._schema == pa.schema([pa.field("""col_1"""       ,   pa.string()   )]   )