





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 							_UpperCamelCase					(	lowerCAmelCase_     ):

       def lowercase (		self: List[Any]		)       -> Tuple:




              """simple docstring"""


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

       def lowercase (		self: str		)       -> List[str]:




              """simple docstring"""


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

       def lowercase (		self: Dict		)       -> Union[str, Any]:




              """simple docstring"""


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

       def lowercase (		self: Union[str, Any]		)       -> Tuple:




              """simple docstring"""


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

       def lowercase (		self: List[str]		)       -> Dict:




              """simple docstring"""


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

       def lowercase (		self: List[str]		)       -> List[Any]:




              """simple docstring"""


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

       def lowercase (		self: List[str]		)       -> int:




              """simple docstring"""


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

       def lowercase (		self: List[Any]		)       -> Union[str, Any]:




              """simple docstring"""


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

       def lowercase (		self: List[Any]		)       -> Tuple:




              """simple docstring"""


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

       def lowercase (		self: Dict		)       -> List[str]:




              """simple docstring"""


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

       def lowercase (		self: List[str]		)       -> Tuple:




              """simple docstring"""


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



       @require_pil
       def lowercase (		self: Tuple		)       -> List[str]:




              """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=_SCREAMING_SNAKE_CASE		) 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"					,   _SCREAMING_SNAKE_CASE		)
                     self.assertFalse(kwargs["optimize_list_casting"]		)




def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Any:
       UpperCamelCase_          =  pa.BufferReader(UpperCamelCase_  ) if isinstance(UpperCamelCase_					,		pa.Buffer  ) else pa.memory_map(UpperCamelCase_  )
       UpperCamelCase_          =  pa.ipc.open_stream(UpperCamelCase_  )
       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, 10]  )
@pytest.mark.parametrize(
    "fields"					,		[None, {"col_1": pa.string(), "col_2": pa.intaa()}, {"col_1": pa.string(), "col_2": pa.intaa()}]  )
def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		int:
       UpperCamelCase_          =  pa.BufferOutputStream()
       UpperCamelCase_          =  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}  )
              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(UpperCamelCase_					,		metadata=writer._schema.metadata  )
       _check_output(output.getvalue()					,		expected_num_chunks=num_examples if writer_batch_size == 1 else 1  )




def  lowerCAmelCase_ (							)     ->		str:
       UpperCamelCase_          =  pa.BufferOutputStream()
       UpperCamelCase_          =  Features({"labels": ClassLabel(names=["neg", "pos"]  )}  )
       with ArrowWriter(stream=UpperCamelCase_					,		features=UpperCamelCase_  ) 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(UpperCamelCase_  )
       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(UpperCamelCase_  )




@pytest.mark.parametrize("writer_batch_size"					,		[None, 1, 10]  )
def  lowerCAmelCase_ (							UpperCamelCase_  )     ->		Any:
       UpperCamelCase_          =  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]  )
                     UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()




@pytest.mark.parametrize("writer_batch_size"					,		[None, 2, 10]  )
def  lowerCAmelCase_ (							UpperCamelCase_  )     ->		str:
       UpperCamelCase_          =  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  )
                     UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()




@pytest.mark.parametrize("writer_batch_size"					,		[None, 2, 10]  )
def  lowerCAmelCase_ (							UpperCamelCase_  )     ->		Any:
       UpperCamelCase_          =  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  )
              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, 10]  )
@pytest.mark.parametrize(
    "fields"					,		[None, {"col_1": pa.string(), "col_2": pa.intaa()}, {"col_1": pa.string(), "col_2": pa.intaa()}]  )
def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Optional[int]:
       UpperCamelCase_          =  pa.BufferOutputStream()
       UpperCamelCase_          =  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": []}  )
              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(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  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Any:
       UpperCamelCase_          =  pa.BufferOutputStream()
       UpperCamelCase_          =  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]}  )  )
              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(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  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Optional[int]:
       UpperCamelCase_          =  pa.BufferOutputStream()
       UpperCamelCase_          =  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]}  )  )
              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(UpperCamelCase_					,		metadata=writer._schema.metadata  )
       _check_output(output.getvalue()					,		expected_num_chunks=num_examples if writer_batch_size == 1 else 1  )




def  lowerCAmelCase_ (							)     ->		Any:
       with tempfile.TemporaryDirectory() as tmp_dir:
              UpperCamelCase_          =  {"col_1": pa.string(), "col_2": pa.intaa()}
              UpperCamelCase_          =  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]}  )
                     UpperCamelCase_       , UpperCamelCase_          =  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  lowerCAmelCase_ (							UpperCamelCase_  )     ->		Dict:
       if pa.types.is_list(UpperCamelCase_  ):
              return get_base_dtype(arr_type.value_type  )
       else:
              return arr_type




def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		List[str]:
       if isinstance(lst[0]					,		UpperCamelCase_  ):
              change_first_primitive_element_in_list(lst[0]					,		UpperCamelCase_  )
       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  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_					,		UpperCamelCase_  )     ->		Dict:
       UpperCamelCase_          =  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  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_					,		UpperCamelCase_  )     ->		Union[str, Any]:
       # in range
       UpperCamelCase_          =  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
              UpperCamelCase_          =  copy.deepcopy(UpperCamelCase_  )
              UpperCamelCase_          =  np.iinfo(expected_dtype.to_pandas_dtype()  ).max + 1
              change_first_primitive_element_in_list(UpperCamelCase_					,		UpperCamelCase_  )
              UpperCamelCase_          =  pa.array(OptimizedTypedSequence(UpperCamelCase_					,		col=UpperCamelCase_  )  )
              assert get_base_dtype(arr.type  ) == pa.intaa()




@pytest.mark.parametrize("raise_exception"					,		[False, True]  )
def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Optional[Any]:
       UpperCamelCase_          =  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  lowerCAmelCase_ (							UpperCamelCase_  )     ->		Optional[int]:
       UpperCamelCase_          =  "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}  )
              UpperCamelCase_       , UpperCamelCase_          =  writer.finalize()
       assert num_examples == 2
       assert num_bytes > 0
       assert mockfs.exists(UpperCamelCase_  )




def  lowerCAmelCase_ (							)     ->		Optional[int]:
       UpperCamelCase_          =  pa.BufferOutputStream()
       with ParquetWriter(stream=UpperCamelCase_  ) 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(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  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_  )     ->		Optional[Any]:
       import PIL.Image

       UpperCamelCase_          =  str(tmp_path / "test_image_rgb.jpg"  )
       PIL.Image.fromarray(np.zeros((5, 5)					,		dtype=np.uinta  )  ).save(UpperCamelCase_					,		format="png"  )
       UpperCamelCase_          =  pa.BufferOutputStream()
       with ParquetWriter(
           stream=UpperCamelCase_					,		features=Features({"image": Image()}  )					,		embed_local_files=UpperCamelCase_  ) as writer:
              writer.write({"image": image_path}  )
              writer.finalize()
       UpperCamelCase_          =  pa.BufferReader(output.getvalue()  )
       UpperCamelCase_          =  pq.read_table(UpperCamelCase_  )
       UpperCamelCase_          =  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  lowerCAmelCase_ (							)     ->		int:
       UpperCamelCase_          =  pa.schema([pa.field("col_1"					,		pa.string()					,		nullable=UpperCamelCase_  )]  )
       UpperCamelCase_          =  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()  )]  )








from ...configuration_utils import PretrainedConfig
from ...utils import logging


_UpperCAmelCase        =				logging.get_logger(__name__)

_UpperCAmelCase        =				{'openai-gpt': 'https://huggingface.co/openai-gpt/resolve/main/config.json'}



class 							_UpperCamelCase					(	lowerCAmelCase_     ):
       _UpperCamelCase				:					List[str]         =   '''openai-gpt'''
       _UpperCamelCase				:					List[str]         =   {
           '''max_position_embeddings''': '''n_positions''',
           '''hidden_size''': '''n_embd''',
           '''num_attention_heads''': '''n_head''',
           '''num_hidden_layers''': '''n_layer''',
       }

       def __init__(		self: int					,   _SCREAMING_SNAKE_CASE: int=40478					,   _SCREAMING_SNAKE_CASE: Optional[Any]=512					,   _SCREAMING_SNAKE_CASE: List[Any]=768					,   _SCREAMING_SNAKE_CASE: str=12					,   _SCREAMING_SNAKE_CASE: Any=12					,   _SCREAMING_SNAKE_CASE: Optional[int]="gelu"					,   _SCREAMING_SNAKE_CASE: int=0.1					,   _SCREAMING_SNAKE_CASE: Optional[int]=0.1					,   _SCREAMING_SNAKE_CASE: Union[str, Any]=0.1					,   _SCREAMING_SNAKE_CASE: Any=1e-5					,   _SCREAMING_SNAKE_CASE: Tuple=0.02					,   _SCREAMING_SNAKE_CASE: str="cls_index"					,   _SCREAMING_SNAKE_CASE: Union[str, Any]=True					,   _SCREAMING_SNAKE_CASE: List[Any]=None					,   _SCREAMING_SNAKE_CASE: List[Any]=True					,   _SCREAMING_SNAKE_CASE: List[Any]=0.1					,   **_SCREAMING_SNAKE_CASE: List[Any]					,   )       -> int:




              """simple docstring"""


              UpperCamelCase_          =  vocab_size
              UpperCamelCase_          =  n_positions
              UpperCamelCase_          =  n_embd
              UpperCamelCase_          =  n_layer
              UpperCamelCase_          =  n_head
              UpperCamelCase_          =  afn
              UpperCamelCase_          =  resid_pdrop
              UpperCamelCase_          =  embd_pdrop
              UpperCamelCase_          =  attn_pdrop
              UpperCamelCase_          =  layer_norm_epsilon
              UpperCamelCase_          =  initializer_range
              UpperCamelCase_          =  summary_type
              UpperCamelCase_          =  summary_use_proj
              UpperCamelCase_          =  summary_activation
              UpperCamelCase_          =  summary_first_dropout
              UpperCamelCase_          =  summary_proj_to_labels
              super().__init__(**_SCREAMING_SNAKE_CASE		)

