import collections
import inspect
import unittest
from typing import Dict, List, Tuple

from transformers import MaskFormerSwinConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, torch_device
from transformers.utils import is_torch_available

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


if is_torch_available():
   import torch
   from torch import nn

   from transformers import MaskFormerSwinBackbone
   from transformers.models.maskformer import MaskFormerSwinModel







class __snake_case       :




   def __init__(					self    :       Union[str, Any] ,			_lowercase    :       List[Any] ,			_lowercase    :       str=13 ,			_lowercase    :       Any=32 ,			_lowercase    :       Union[str, Any]=2 ,			_lowercase    :       Any=3 ,			_lowercase    :       str=16 ,			_lowercase    :       List[Any]=[1, 2, 1] ,			_lowercase    :       Optional[Any]=[2, 2, 4] ,			_lowercase    :       Tuple=2 ,			_lowercase    :       Union[str, Any]=2.0 ,			_lowercase    :       Tuple=True ,			_lowercase    :       str=0.0 ,			_lowercase    :       Optional[Any]=0.0 ,			_lowercase    :       str=0.1 ,			_lowercase    :       Optional[int]="gelu" ,			_lowercase    :       List[str]=False ,			_lowercase    :       Union[str, Any]=True ,			_lowercase    :       Dict=0.02 ,			_lowercase    :       Optional[Any]=1E-5 ,			_lowercase    :       Optional[int]=True ,			_lowercase    :       Union[str, Any]=None ,			_lowercase    :       List[Any]=True ,			_lowercase    :       Tuple=10 ,			_lowercase    :       List[Any]=8 ,			_lowercase    :       List[str]=["stage1", "stage2", "stage3"] ,			_lowercase    :       List[str]=[1, 2, 3] ,			):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   parent
     SCREAMING_SNAKE_CASE__    							=   batch_size
     SCREAMING_SNAKE_CASE__    							=   image_size
     SCREAMING_SNAKE_CASE__    							=   patch_size
     SCREAMING_SNAKE_CASE__    							=   num_channels
     SCREAMING_SNAKE_CASE__    							=   embed_dim
     SCREAMING_SNAKE_CASE__    							=   depths
     SCREAMING_SNAKE_CASE__    							=   num_heads
     SCREAMING_SNAKE_CASE__    							=   window_size
     SCREAMING_SNAKE_CASE__    							=   mlp_ratio
     SCREAMING_SNAKE_CASE__    							=   qkv_bias
     SCREAMING_SNAKE_CASE__    							=   hidden_dropout_prob
     SCREAMING_SNAKE_CASE__    							=   attention_probs_dropout_prob
     SCREAMING_SNAKE_CASE__    							=   drop_path_rate
     SCREAMING_SNAKE_CASE__    							=   hidden_act
     SCREAMING_SNAKE_CASE__    							=   use_absolute_embeddings
     SCREAMING_SNAKE_CASE__    							=   patch_norm
     SCREAMING_SNAKE_CASE__    							=   layer_norm_eps
     SCREAMING_SNAKE_CASE__    							=   initializer_range
     SCREAMING_SNAKE_CASE__    							=   is_training
     SCREAMING_SNAKE_CASE__    							=   scope
     SCREAMING_SNAKE_CASE__    							=   use_labels
     SCREAMING_SNAKE_CASE__    							=   type_sequence_label_size
     SCREAMING_SNAKE_CASE__    							=   encoder_stride
     SCREAMING_SNAKE_CASE__    							=   out_features
     SCREAMING_SNAKE_CASE__    							=   out_indices




   def 			__a		(					self    :       Optional[int]   ):



     """simple docstring"""






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

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

     SCREAMING_SNAKE_CASE__    							=   self.get_config()

     return config, pixel_values, labels




   def 			__a		(					self    :       List[str]   ):



     """simple docstring"""






     return MaskFormerSwinConfig(
         image_size=self.image_size ,			patch_size=self.patch_size ,			num_channels=self.num_channels ,			embed_dim=self.embed_dim ,			depths=self.depths ,			num_heads=self.num_heads ,			window_size=self.window_size ,			mlp_ratio=self.mlp_ratio ,			qkv_bias=self.qkv_bias ,			hidden_dropout_prob=self.hidden_dropout_prob ,			attention_probs_dropout_prob=self.attention_probs_dropout_prob ,			drop_path_rate=self.drop_path_rate ,			hidden_act=self.hidden_act ,			use_absolute_embeddings=self.use_absolute_embeddings ,			path_norm=self.patch_norm ,			layer_norm_eps=self.layer_norm_eps ,			initializer_range=self.initializer_range ,			encoder_stride=self.encoder_stride ,			out_features=self.out_features ,			out_indices=self.out_indices ,			)




   def 			__a		(					self    :       Optional[int] ,			_lowercase    :       int ,			_lowercase    :       List[str] ,			_lowercase    :       Optional[Any]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   MaskFormerSwinModel(config=_lowercase   )
     model.to(_lowercase   )
     model.eval()
     SCREAMING_SNAKE_CASE__    							=   model(_lowercase   )

     SCREAMING_SNAKE_CASE__    							=   ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths   ) - 1))
     SCREAMING_SNAKE_CASE__    							=   int(config.embed_dim * 2 ** (len(config.depths   ) - 1)   )

     self.parent.assertEqual(result.last_hidden_state.shape ,			(self.batch_size, expected_seq_len, expected_dim)   )




   def 			__a		(					self    :       Any ,			_lowercase    :       Optional[Any] ,			_lowercase    :       int ,			_lowercase    :       Optional[Any]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   MaskFormerSwinBackbone(config=_lowercase   )
     model.to(_lowercase   )
     model.eval()
     SCREAMING_SNAKE_CASE__    							=   model(_lowercase   )

     # verify feature maps
     self.parent.assertEqual(len(result.feature_maps   ) ,			len(config.out_features   )   )
     self.parent.assertListEqual(list(result.feature_maps[0].shape   ) ,			[13, 16, 16, 16]   )

     # verify channels
     self.parent.assertEqual(len(model.channels   ) ,			len(config.out_features   )   )
     self.parent.assertListEqual(model.channels ,			[16, 32, 64]   )

     # verify ValueError
     with self.parent.assertRaises(_lowercase   ):
       SCREAMING_SNAKE_CASE__    							=   ["""stem"""]
       SCREAMING_SNAKE_CASE__    							=   MaskFormerSwinBackbone(config=_lowercase   )




   def 			__a		(					self    :       Any   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.prepare_config_and_inputs()
     SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   config_and_inputs
     SCREAMING_SNAKE_CASE__    							=   {"""pixel_values""": pixel_values}
     return config, inputs_dict







@require_torch
class __snake_case       (		lowerCamelCase_     ,	lowerCamelCase_     ,	unittest.TestCase	):
   lowerCAmelCase_														=							(
       (
           MaskFormerSwinModel,
           MaskFormerSwinBackbone,
       )
       if is_torch_available()
       else ()
   )
   lowerCAmelCase_														=							{"feature-extraction": MaskFormerSwinModel} if is_torch_available() else {}
   lowerCAmelCase_														=							False
   lowerCAmelCase_														=							False
   lowerCAmelCase_														=							False
   lowerCAmelCase_														=							False
   lowerCAmelCase_														=							False




   def 			__a		(					self    :       Tuple   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   MaskFormerSwinModelTester(self   )
     SCREAMING_SNAKE_CASE__    							=   ConfigTester(self ,			config_class=_lowercase ,			embed_dim=37   )




   @require_torch_multi_gpu
   @unittest.skip(
       reason=(
           """`MaskFormerSwinModel` outputs `hidden_states_spatial_dimensions` which doesn't work well with"""
           """ `nn.DataParallel`"""
       )   )
   def 			__a		(					self    :       Dict   ):



     """simple docstring"""






     pass




   def 			__a		(					self    :       Any   ):



     """simple docstring"""






     self.create_and_test_config_common_properties()
     self.config_tester.create_and_test_config_to_json_string()
     self.config_tester.create_and_test_config_to_json_file()
     self.config_tester.create_and_test_config_from_and_save_pretrained()
     self.config_tester.create_and_test_config_with_num_labels()
     self.config_tester.check_config_can_be_init_without_params()
     self.config_tester.check_config_arguments_init()




   def 			__a		(					self    :       Any   ):



     """simple docstring"""






     return




   def 			__a		(					self    :       str   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*_lowercase   )




   def 			__a		(					self    :       Dict   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_backbone(*_lowercase   )




   @unittest.skip("""Swin does not use inputs_embeds"""   )
   def 			__a		(					self    :       Any   ):



     """simple docstring"""






     pass




   @unittest.skip("""Swin does not support feedforward chunking"""   )
   def 			__a		(					self    :       Dict   ):



     """simple docstring"""






     pass




   def 			__a		(					self    :       Tuple   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs_for_common()

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




   def 			__a		(					self    :       Tuple   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs_for_common()

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

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




   @unittest.skip(reason="""MaskFormerSwin is only used as backbone and doesn't support output_attentions"""   )
   def 			__a		(					self    :       List[str]   ):



     """simple docstring"""






     pass




   @unittest.skip(reason="""MaskFormerSwin is only used as an internal backbone"""   )
   def 			__a		(					self    :       List[Any]   ):



     """simple docstring"""






     pass




   def 			__a		(					self    :       List[str] ,			_lowercase    :       List[str] ,			_lowercase    :       List[str] ,			_lowercase    :       List[Any] ,			_lowercase    :       Optional[int]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   model_class(_lowercase   )
     model.to(_lowercase   )
     model.eval()

     with torch.no_grad():
       SCREAMING_SNAKE_CASE__    							=   model(**self._prepare_for_class(_lowercase ,			_lowercase   )   )

     SCREAMING_SNAKE_CASE__    							=   outputs.hidden_states

     SCREAMING_SNAKE_CASE__    							=   getattr(
         self.model_tester ,			"""expected_num_hidden_layers""" ,			len(self.model_tester.depths   ) + 1   )
     self.assertEqual(len(_lowercase   ) ,			_lowercase   )

     # Swin has a different seq_length
     SCREAMING_SNAKE_CASE__    							=   (
         config.patch_size
         if isinstance(config.patch_size ,			collections.abc.Iterable   )
         else (config.patch_size, config.patch_size)
     )

     SCREAMING_SNAKE_CASE__    							=   (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

     self.assertListEqual(
         list(hidden_states[0].shape[-2:]   ) ,			[num_patches, self.model_tester.embed_dim] ,			)




   def 			__a		(					self    :       int   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs_for_common()

     SCREAMING_SNAKE_CASE__    							=   (
         self.model_tester.image_size
         if isinstance(self.model_tester.image_size ,			collections.abc.Iterable   )
         else (self.model_tester.image_size, self.model_tester.image_size)
     )

     for model_class in self.all_model_classes:
       SCREAMING_SNAKE_CASE__    							=   True
       self.check_hidden_states_output(_lowercase ,			_lowercase ,			_lowercase ,			_lowercase   )

       # check that output_hidden_states also work using config
       del inputs_dict["output_hidden_states"]
       SCREAMING_SNAKE_CASE__    							=   True

       self.check_hidden_states_output(_lowercase ,			_lowercase ,			_lowercase ,			_lowercase   )




   def 			__a		(					self    :       Dict   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE__    							=   3

     SCREAMING_SNAKE_CASE__    							=   (
         self.model_tester.image_size
         if isinstance(self.model_tester.image_size ,			collections.abc.Iterable   )
         else (self.model_tester.image_size, self.model_tester.image_size)
     )
     SCREAMING_SNAKE_CASE__    							=   (
         config.patch_size
         if isinstance(config.patch_size ,			collections.abc.Iterable   )
         else (config.patch_size, config.patch_size)
     )

     SCREAMING_SNAKE_CASE__    							=   image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
     SCREAMING_SNAKE_CASE__    							=   image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

     for model_class in self.all_model_classes:
       SCREAMING_SNAKE_CASE__    							=   True
       self.check_hidden_states_output(_lowercase ,			_lowercase ,			_lowercase ,			(padded_height, padded_width)   )

       # check that output_hidden_states also work using config
       del inputs_dict["output_hidden_states"]
       SCREAMING_SNAKE_CASE__    							=   True
       self.check_hidden_states_output(_lowercase ,			_lowercase ,			_lowercase ,			(padded_height, padded_width)   )




   @unittest.skip(reason="""MaskFormerSwin doesn't have pretrained checkpoints"""   )
   def 			__a		(					self    :       str   ):



     """simple docstring"""






     pass




   @unittest.skip(reason="""This will be fixed once MaskFormerSwin is replaced by native Swin"""   )
   def 			__a		(					self    :       List[Any]   ):



     """simple docstring"""






     pass




   @unittest.skip(reason="""This will be fixed once MaskFormerSwin is replaced by native Swin"""   )
   def 			__a		(					self    :       List[str]   ):



     """simple docstring"""






     pass




   def 			__a		(					self    :       Optional[Any]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs_for_common()

     def set_nan_tensor_to_zero(_lowercase    :       List[Any]   ):
       SCREAMING_SNAKE_CASE__    							=   0
       return t

     def check_equivalence(_lowercase    :       List[str] ,			_lowercase    :       Union[str, Any] ,			_lowercase    :       Optional[Any] ,			_lowercase    :       int={}   ):
       with torch.no_grad():
         SCREAMING_SNAKE_CASE__    							=   model(**_lowercase ,			return_dict=_lowercase ,			**_lowercase   )
         SCREAMING_SNAKE_CASE__    							=   model(**_lowercase ,			return_dict=_lowercase ,			**_lowercase   ).to_tuple()

         def recursive_check(_lowercase    :       List[Any] ,			_lowercase    :       Any   ):
           if isinstance(_lowercase ,			(List, Tuple)   ):
             for tuple_iterable_value, dict_iterable_value in zip(_lowercase ,			_lowercase   ):
               recursive_check(_lowercase ,			_lowercase   )
           elif isinstance(_lowercase ,			_lowercase   ):
             for tuple_iterable_value, dict_iterable_value in zip(
                 tuple_object.values() ,			dict_object.values()   ):
               recursive_check(_lowercase ,			_lowercase   )
           elif tuple_object is None:
             return
           else:
             self.assertTrue(
                 torch.allclose(
                     set_nan_tensor_to_zero(_lowercase   ) ,			set_nan_tensor_to_zero(_lowercase   ) ,			atol=1E-5   ) ,			msg=(
                     """Tuple and dict output are not equal. Difference:"""
                     f""" {torch.max(torch.abs(tuple_object - dict_object   )   )}. Tuple has `nan`:"""
                     f""" {torch.isnan(_lowercase   ).any()} and `inf`: {torch.isinf(_lowercase   )}. Dict has"""
                     f""" `nan`: {torch.isnan(_lowercase   ).any()} and `inf`: {torch.isinf(_lowercase   )}."""
                 ) ,			)

         recursive_check(_lowercase ,			_lowercase   )

     for model_class in self.all_model_classes:
       SCREAMING_SNAKE_CASE__    							=   model_class(_lowercase   )
       model.to(_lowercase   )
       model.eval()

       SCREAMING_SNAKE_CASE__    							=   self._prepare_for_class(_lowercase ,			_lowercase   )
       SCREAMING_SNAKE_CASE__    							=   self._prepare_for_class(_lowercase ,			_lowercase   )
       check_equivalence(_lowercase ,			_lowercase ,			_lowercase   )

       SCREAMING_SNAKE_CASE__    							=   self._prepare_for_class(_lowercase ,			_lowercase ,			return_labels=_lowercase   )
       SCREAMING_SNAKE_CASE__    							=   self._prepare_for_class(_lowercase ,			_lowercase ,			return_labels=_lowercase   )
       check_equivalence(_lowercase ,			_lowercase ,			_lowercase   )

       SCREAMING_SNAKE_CASE__    							=   self._prepare_for_class(_lowercase ,			_lowercase   )
       SCREAMING_SNAKE_CASE__    							=   self._prepare_for_class(_lowercase ,			_lowercase   )
       check_equivalence(_lowercase ,			_lowercase ,			_lowercase ,			{"""output_hidden_states""": True}   )

       SCREAMING_SNAKE_CASE__    							=   self._prepare_for_class(_lowercase ,			_lowercase ,			return_labels=_lowercase   )
       SCREAMING_SNAKE_CASE__    							=   self._prepare_for_class(_lowercase ,			_lowercase ,			return_labels=_lowercase   )
       check_equivalence(_lowercase ,			_lowercase ,			_lowercase ,			{"""output_hidden_states""": True}   )







@require_torch
class __snake_case       (		unittest.TestCase     ,	lowerCamelCase_	):
   lowerCAmelCase_														=							(MaskFormerSwinBackbone,) if is_torch_available() else ()
   lowerCAmelCase_														=							MaskFormerSwinConfig




   def 			__a		(					self    :       Optional[int]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   MaskFormerSwinModelTester(self   )




   def 			__a		(					self    :       Optional[Any]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs_for_common()
     SCREAMING_SNAKE_CASE__    							=   inputs_dict["""pixel_values"""].shape[0]

     for backbone_class in self.all_model_classes:
       SCREAMING_SNAKE_CASE__    							=   backbone_class(_lowercase   )
       backbone.to(_lowercase   )
       backbone.eval()

       SCREAMING_SNAKE_CASE__    							=   backbone(**_lowercase   )

       # Test default outputs and verify feature maps
       self.assertIsInstance(outputs.feature_maps ,			_lowercase   )
       self.assertTrue(len(outputs.feature_maps   ) == len(backbone.channels   )   )
       for feature_map, n_channels in zip(outputs.feature_maps ,			backbone.channels   ):
         self.assertTrue(feature_map.shape[:2] ,			(batch_size, n_channels)   )
       self.assertIsNone(outputs.hidden_states   )
       self.assertIsNone(outputs.attentions   )

       # Test output_hidden_states=True
       SCREAMING_SNAKE_CASE__    							=   backbone(**_lowercase ,			output_hidden_states=_lowercase   )
       self.assertIsNotNone(outputs.hidden_states   )
       self.assertTrue(len(outputs.hidden_states   ) ,			len(backbone.stage_names   )   )
       # We skip the stem layer
       for hidden_states, n_channels in zip(outputs.hidden_states[1:] ,			backbone.channels   ):
         for hidden_state in hidden_states:
           # Hidden states are in the format (batch_size, (height * width), n_channels)
           SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   hidden_state.shape
           self.assertTrue((h_batch_size, h_n_channels) ,			(batch_size, n_channels)   )

            # Test output_attentions=True
       if self.has_attentions:
         SCREAMING_SNAKE_CASE__    							=   backbone(**_lowercase ,			output_attentions=_lowercase   )
         self.assertIsNotNone(outputs.attentions   )

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 __snake_case       (		lowerCamelCase_	):




   def 			__a		(					self    :       Optional[Any]   ):



     """simple docstring"""






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




   def 			__a		(					self    :       Tuple   ):



     """simple docstring"""






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




   def 			__a		(					self    :       int   ):



     """simple docstring"""






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




   def 			__a		(					self    :       List[str]   ):



     """simple docstring"""






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




   def 			__a		(					self    :       Optional[Any]   ):



     """simple docstring"""






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




   def 			__a		(					self    :       str   ):



     """simple docstring"""






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




   def 			__a		(					self    :       Any   ):



     """simple docstring"""






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




   def 			__a		(					self    :       Tuple   ):



     """simple docstring"""






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




   def 			__a		(					self    :       Dict   ):



     """simple docstring"""






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




   def 			__a		(					self    :       Optional[int]   ):



     """simple docstring"""






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




   def 			__a		(					self    :       int   ):



     """simple docstring"""






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




   @require_pil
   def 			__a		(					self    :       List[str]   ):



     """simple docstring"""






     import PIL.Image

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





def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Optional[Any]				,  __UpperCamelCase					: int			)							->      str:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__    							=   pa.BufferReader(__UpperCamelCase			) if isinstance(__UpperCamelCase				,  pa.Buffer			) else pa.memory_map(__UpperCamelCase			)
  SCREAMING_SNAKE_CASE__    							=   pa.ipc.open_stream(__UpperCamelCase			)
  SCREAMING_SNAKE_CASE__    							=   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, 10]			)
@pytest.mark.parametrize(
    """fields"""				,  [None, {"""col_1""": pa.string(), """col_2""": pa.intaa()}, {"""col_1""": pa.string(), """col_2""": pa.intaa()}]			)
def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Union[str, Any]				,  __UpperCamelCase					: int			)							->      List[str]:







  """simple docstring"""





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



def 			__SCREAMING_SNAKE_CASE						(      )							->      int:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__    							=   pa.BufferOutputStream()
  SCREAMING_SNAKE_CASE__    							=   Features({"""labels""": ClassLabel(names=["""neg""", """pos"""]			)}			)
  with ArrowWriter(stream=__UpperCamelCase				,  features=__UpperCamelCase			) as writer:
    writer.write({"""labels""": 0}			)
    writer.write({"""labels""": 1}			)
    SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   writer.finalize()
  assert num_examples == 2
  assert num_bytes > 0
  assert writer._schema == features.arrow_schema
  assert writer._schema.metadata == features.arrow_schema.metadata
  SCREAMING_SNAKE_CASE__    							=   pa.BufferReader(output.getvalue()			)
  SCREAMING_SNAKE_CASE__    							=   pa.ipc.open_stream(__UpperCamelCase			)
  SCREAMING_SNAKE_CASE__    							=   f.read_all()
  SCREAMING_SNAKE_CASE__    							=   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(__UpperCamelCase			)



@pytest.mark.parametrize("""writer_batch_size"""				,  [None, 1, 10]			)
def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: List[str]			)							->      List[str]:







  """simple docstring"""





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



@pytest.mark.parametrize("""writer_batch_size"""				,  [None, 2, 10]			)
def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Dict			)							->      Dict:







  """simple docstring"""





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



@pytest.mark.parametrize("""writer_batch_size"""				,  [None, 2, 10]			)
def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: List[str]			)							->      int:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__    							=   pa.BufferOutputStream()
  with ArrowWriter(
      stream=__UpperCamelCase				,  writer_batch_size=__UpperCamelCase				,  hash_salt="""split_name"""				,  check_duplicates=__UpperCamelCase				,  ) as writer:
    writer.write({"""col_1""": """foo""", """col_2""": 1}				,  key=1			)
    writer.write({"""col_1""": """bar""", """col_2""": 2}				,  key=2			)
    SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   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, 10]			)
@pytest.mark.parametrize(
    """fields"""				,  [None, {"""col_1""": pa.string(), """col_2""": pa.intaa()}, {"""col_1""": pa.string(), """col_2""": pa.intaa()}]			)
def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Optional[int]				,  __UpperCamelCase					: Union[str, Any]			)							->      Optional[int]:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__    							=   pa.BufferOutputStream()
  SCREAMING_SNAKE_CASE__    							=   pa.schema(__UpperCamelCase			) if fields else None
  with ArrowWriter(stream=__UpperCamelCase				,  schema=__UpperCamelCase				,  writer_batch_size=__UpperCamelCase			) as writer:
    writer.write_batch({"""col_1""": ["""foo""", """bar"""], """col_2""": [1, 2]}			)
    writer.write_batch({"""col_1""": [], """col_2""": []}			)
    SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   writer.finalize()
  assert num_examples == 2
  assert num_bytes > 0
  if not fields:
    SCREAMING_SNAKE_CASE__    							=   {"""col_1""": pa.string(), """col_2""": pa.intaa()}
  assert writer._schema == pa.schema(__UpperCamelCase				,  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, 10]			)
@pytest.mark.parametrize(
    """fields"""				,  [None, {"""col_1""": pa.string(), """col_2""": pa.intaa()}, {"""col_1""": pa.string(), """col_2""": pa.intaa()}]			)
def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: int				,  __UpperCamelCase					: Union[str, Any]			)							->      Any:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__    							=   pa.BufferOutputStream()
  SCREAMING_SNAKE_CASE__    							=   pa.schema(__UpperCamelCase			) if fields else None
  with ArrowWriter(stream=__UpperCamelCase				,  schema=__UpperCamelCase				,  writer_batch_size=__UpperCamelCase			) as writer:
    writer.write_table(pa.Table.from_pydict({"""col_1""": ["""foo""", """bar"""], """col_2""": [1, 2]}			)			)
    SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   writer.finalize()
  assert num_examples == 2
  assert num_bytes > 0
  if not fields:
    SCREAMING_SNAKE_CASE__    							=   {"""col_1""": pa.string(), """col_2""": pa.intaa()}
  assert writer._schema == pa.schema(__UpperCamelCase				,  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, 10]			)
@pytest.mark.parametrize(
    """fields"""				,  [None, {"""col_1""": pa.string(), """col_2""": pa.intaa()}, {"""col_1""": pa.string(), """col_2""": pa.intaa()}]			)
def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Union[str, Any]				,  __UpperCamelCase					: List[str]			)							->      str:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__    							=   pa.BufferOutputStream()
  SCREAMING_SNAKE_CASE__    							=   pa.schema(__UpperCamelCase			) if fields else None
  with ArrowWriter(stream=__UpperCamelCase				,  schema=__UpperCamelCase				,  writer_batch_size=__UpperCamelCase			) 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]}			)			)
    SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   writer.finalize()
  assert num_examples == 2
  assert num_bytes > 0
  if not fields:
    SCREAMING_SNAKE_CASE__    							=   {"""col_1""": pa.string(), """col_2""": pa.intaa()}
  assert writer._schema == pa.schema(__UpperCamelCase				,  metadata=writer._schema.metadata			)
  _check_output(output.getvalue()				,  expected_num_chunks=num_examples if writer_batch_size == 1 else 1			)



def 			__SCREAMING_SNAKE_CASE						(      )							->      List[Any]:







  """simple docstring"""





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



def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: int			)							->      Tuple:







  """simple docstring"""





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



def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Dict				,  __UpperCamelCase					: Union[str, Any]			)							->      Tuple:







  """simple docstring"""





  if isinstance(lst[0]				,  __UpperCamelCase			):
    change_first_primitive_element_in_list(lst[0]				,  __UpperCamelCase			)
  else:
    SCREAMING_SNAKE_CASE__    							=   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 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Union[str, Any]				,  __UpperCamelCase					: Optional[int]				,  __UpperCamelCase					: Dict			)							->      Any:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__    							=   pa.array(TypedSequence(__UpperCamelCase				,  optimized_int_type=__UpperCamelCase			)			)
  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 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: List[Any]				,  __UpperCamelCase					: Dict				,  __UpperCamelCase					: List[str]			)							->      Optional[int]:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__    							=   pa.array(OptimizedTypedSequence(__UpperCamelCase				,  col=__UpperCamelCase			)			)
  assert get_base_dtype(arr.type			) == expected_dtype

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



@pytest.mark.parametrize("""raise_exception"""				,  [False, True]			)
def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Optional[int]				,  __UpperCamelCase					: Union[str, Any]			)							->      str:







  """simple docstring"""





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



def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Tuple			)							->      str:







  """simple docstring"""





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



def 			__SCREAMING_SNAKE_CASE						(      )							->      str:







  """simple docstring"""





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



@require_pil
@pytest.mark.parametrize("""embed_local_files"""				,  [False, True]			)
def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Tuple				,  __UpperCamelCase					: Tuple			)							->      Optional[Any]:







  """simple docstring"""





  import PIL.Image

  SCREAMING_SNAKE_CASE__    							=   str(tmp_path / """test_image_rgb.jpg"""			)
  PIL.Image.fromarray(np.zeros((5, 5)				,  dtype=np.uinta			)			).save(__UpperCamelCase				,  format="""png"""			)
  SCREAMING_SNAKE_CASE__    							=   pa.BufferOutputStream()
  with ParquetWriter(
      stream=__UpperCamelCase				,  features=Features({"""image""": Image()}			)				,  embed_local_files=__UpperCamelCase			) as writer:
    writer.write({"""image""": image_path}			)
    writer.finalize()
  SCREAMING_SNAKE_CASE__    							=   pa.BufferReader(output.getvalue()			)
  SCREAMING_SNAKE_CASE__    							=   pq.read_table(__UpperCamelCase			)
  SCREAMING_SNAKE_CASE__    							=   pa_table.to_pydict()
  if embed_local_files:
    assert isinstance(out["""image"""][0]["""path"""]				,  __UpperCamelCase			)
    with open(__UpperCamelCase				,  """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 			__SCREAMING_SNAKE_CASE						(      )							->      List[Any]:







  """simple docstring"""





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