




import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, BatchEncoding, PLBartTokenizer, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
)

from ...test_tokenization_common import TokenizerTesterMixin


lowerCamelCase__					:  List[str]      							=				get_tests_dir('fixtures/test_sentencepiece.model')


if is_torch_available():
   from transformers.models.plbart.modeling_plbart import shift_tokens_right

lowerCamelCase__					:  int      							=				50_003
lowerCamelCase__					:  List[str]      							=				50_002





@require_sentencepiece
@require_tokenizers
class 			lowerCamelCase_				( _SCREAMING_SNAKE_CASE			,  unittest.TestCase   ):


 '''simple docstring'''







 lowercase_							       =       PLBartTokenizer
 lowercase_							       =       None
 lowercase_							       =       False



 def  lowerCAmelCase_    (				self    :						List[str]		):
       super().setUp()

       # We have a SentencePiece fixture for testing
       SCREAMING_SNAKE_CASE_      		=     PLBartTokenizer(_lowerCAmelCase							,							language_codes='base'							,							keep_accents=_lowerCAmelCase		)
       tokenizer.save_pretrained(self.tmpdirname		)



 def  lowerCAmelCase_    (				self    :						int		):
       SCREAMING_SNAKE_CASE_      		=     PLBartTokenizer(_lowerCAmelCase							,							language_codes='base'							,							keep_accents=_lowerCAmelCase		)

       SCREAMING_SNAKE_CASE_      		=     tokenizer.tokenize('This is a test'		)
       self.assertListEqual(_lowerCAmelCase							,							['▁This', '▁is', '▁a', '▁t', 'est']		)

       self.assertListEqual(
           tokenizer.convert_tokens_to_ids(_lowerCAmelCase		)							,							[value + tokenizer.fairseq_offset for value in [285, 46, 10, 170, 382]]							,							)

       SCREAMING_SNAKE_CASE_      		=     tokenizer.tokenize('I was born in 92000, and this is falsé.'		)
       self.assertListEqual(
           _lowerCAmelCase							,							[
               SPIECE_UNDERLINE + 'I',
               SPIECE_UNDERLINE + 'was',
               SPIECE_UNDERLINE + 'b',
               'or',
               'n',
               SPIECE_UNDERLINE + 'in',
               SPIECE_UNDERLINE + '',
               '9',
               '2',
               '0',
               '0',
               '0',
               ',',
               SPIECE_UNDERLINE + 'and',
               SPIECE_UNDERLINE + 'this',
               SPIECE_UNDERLINE + 'is',
               SPIECE_UNDERLINE + 'f',
               'al',
               's',
               'é',
               '.',
           ]							,							)
       SCREAMING_SNAKE_CASE_      		=     tokenizer.convert_tokens_to_ids(_lowerCAmelCase		)
       self.assertListEqual(
           _lowerCAmelCase							,							[
               value + tokenizer.fairseq_offset
               for value in [8, 21, 84, 55, 24, 19, 7, 2, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 2, 4]
           ]							,							)

       SCREAMING_SNAKE_CASE_      		=     tokenizer.convert_ids_to_tokens(_lowerCAmelCase		)
       self.assertListEqual(
           _lowerCAmelCase							,							[
               SPIECE_UNDERLINE + 'I',
               SPIECE_UNDERLINE + 'was',
               SPIECE_UNDERLINE + 'b',
               'or',
               'n',
               SPIECE_UNDERLINE + 'in',
               SPIECE_UNDERLINE + '',
               '<unk>',
               '2',
               '0',
               '0',
               '0',
               ',',
               SPIECE_UNDERLINE + 'and',
               SPIECE_UNDERLINE + 'this',
               SPIECE_UNDERLINE + 'is',
               SPIECE_UNDERLINE + 'f',
               'al',
               's',
               '<unk>',
               '.',
           ]							,							)

       SCREAMING_SNAKE_CASE_      		=     tokenizer.vocab_size
       SCREAMING_SNAKE_CASE_      		=     [tokenizer.convert_ids_to_tokens(_lowerCAmelCase		) for x in range(end - 4							,							_lowerCAmelCase		)]

       self.assertListEqual(_lowerCAmelCase							,							['__java__', '__python__', '__en_XX__', '<mask>']		)

       SCREAMING_SNAKE_CASE_      		=     'java.lang.Exception, python.lang.Exception, javascript, php, ruby, go'
       SCREAMING_SNAKE_CASE_      		=     tokenizer(_lowerCAmelCase		).input_ids
       self.assertEqual(
           tokenizer.decode(_lowerCAmelCase							,							skip_special_tokens=_lowerCAmelCase							,							clean_up_tokenization_spaces=_lowerCAmelCase		)							,							_lowerCAmelCase							,							)




 def  lowerCAmelCase_    (				self    :						int		):
       SCREAMING_SNAKE_CASE_      		=     PLBartTokenizer(_lowerCAmelCase							,							language_codes='multi'							,							keep_accents=_lowerCAmelCase		)

       SCREAMING_SNAKE_CASE_      		=     tokenizer.tokenize('This is a test'		)
       self.assertListEqual(_lowerCAmelCase							,							['▁This', '▁is', '▁a', '▁t', 'est']		)

       self.assertListEqual(
           tokenizer.convert_tokens_to_ids(_lowerCAmelCase		)							,							[value + tokenizer.fairseq_offset for value in [285, 46, 10, 170, 382]]							,							)

       SCREAMING_SNAKE_CASE_      		=     tokenizer.tokenize('I was born in 92000, and this is falsé.'		)
       self.assertListEqual(
           _lowerCAmelCase							,							[
               SPIECE_UNDERLINE + 'I',
               SPIECE_UNDERLINE + 'was',
               SPIECE_UNDERLINE + 'b',
               'or',
               'n',
               SPIECE_UNDERLINE + 'in',
               SPIECE_UNDERLINE + '',
               '9',
               '2',
               '0',
               '0',
               '0',
               ',',
               SPIECE_UNDERLINE + 'and',
               SPIECE_UNDERLINE + 'this',
               SPIECE_UNDERLINE + 'is',
               SPIECE_UNDERLINE + 'f',
               'al',
               's',
               'é',
               '.',
           ]							,							)
       SCREAMING_SNAKE_CASE_      		=     tokenizer.convert_tokens_to_ids(_lowerCAmelCase		)
       self.assertListEqual(
           _lowerCAmelCase							,							[
               value + tokenizer.fairseq_offset
               for value in [8, 21, 84, 55, 24, 19, 7, 2, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 2, 4]
           ]							,							)

       SCREAMING_SNAKE_CASE_      		=     tokenizer.convert_ids_to_tokens(_lowerCAmelCase		)
       self.assertListEqual(
           _lowerCAmelCase							,							[
               SPIECE_UNDERLINE + 'I',
               SPIECE_UNDERLINE + 'was',
               SPIECE_UNDERLINE + 'b',
               'or',
               'n',
               SPIECE_UNDERLINE + 'in',
               SPIECE_UNDERLINE + '',
               '<unk>',
               '2',
               '0',
               '0',
               '0',
               ',',
               SPIECE_UNDERLINE + 'and',
               SPIECE_UNDERLINE + 'this',
               SPIECE_UNDERLINE + 'is',
               SPIECE_UNDERLINE + 'f',
               'al',
               's',
               '<unk>',
               '.',
           ]							,							)
       SCREAMING_SNAKE_CASE_      		=     tokenizer.vocab_size
       SCREAMING_SNAKE_CASE_      		=     [tokenizer.convert_ids_to_tokens(_lowerCAmelCase		) for x in range(end - 7							,							_lowerCAmelCase		)]

       self.assertListEqual(
           _lowerCAmelCase							,							['__java__', '__python__', '__en_XX__', '__javascript__', '__php__', '__ruby__', '__go__']		)
       SCREAMING_SNAKE_CASE_      		=     'java.lang.Exception, python.lang.Exception, javascript, php, ruby, go'
       SCREAMING_SNAKE_CASE_      		=     tokenizer(_lowerCAmelCase		).input_ids
       self.assertEqual(
           tokenizer.decode(_lowerCAmelCase							,							skip_special_tokens=_lowerCAmelCase							,							clean_up_tokenization_spaces=_lowerCAmelCase		)							,							_lowerCAmelCase							,							)





@require_torch
@require_sentencepiece
@require_tokenizers
class 			lowerCamelCase_				( unittest.TestCase   ):


 '''simple docstring'''







 lowercase_							       =       "uclanlp/plbart-python-en_XX"
 lowercase_							       =       [
     "def maximum(a,b,c):NEW_LINE_INDENTreturn max([a,b,c])",
     "def sum(a,b,c):NEW_LINE_INDENTreturn sum([a,b,c])",
 ]
 lowercase_							       =       [
     "Returns the maximum value of a b c.",
     "Sums the values of a b c.",
 ]
 lowercase_							       =       [
     134,
     5_452,
     33_460,
     33_441,
     33_463,
     33_465,
     33_463,
     33_449,
     988,
     20,
     33_456,
     19,
     33_456,
     771,
     39,
     4_258,
     889,
     3_318,
     33_441,
     33_463,
     33_465,
     33_463,
     33_449,
     2_471,
     2,
     PYTHON_CODE,
 ]



 @classmethod
 def  lowerCAmelCase_    (				cls    :						Optional[int]		):
       SCREAMING_SNAKE_CASE_      		=     PLBartTokenizer.from_pretrained(
           cls.checkpoint_name							,							language_codes='base'							,							src_lang='python'							,							tgt_lang='en_XX'		)
       SCREAMING_SNAKE_CASE_      		=     1
       return cls



 def  lowerCAmelCase_    (				self    :						Union[str, Any]		):
       self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['__java__']							,							50_001		)
       self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['__python__']							,							50_002		)
       self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['__en_XX__']							,							50_003		)



 def  lowerCAmelCase_    (				self    :						Any		):
       SCREAMING_SNAKE_CASE_      		=     self.tokenizer.batch_encode_plus(self.src_text		).input_ids[0]
       self.assertListEqual(self.expected_src_tokens							,							_lowerCAmelCase		)



 def  lowerCAmelCase_    (				self    :						Any		):
       self.assertIn(_lowerCAmelCase							,							self.tokenizer.all_special_ids		)
       SCREAMING_SNAKE_CASE_      		=     [EN_CODE, 9_037, 33_442, 57, 752, 153, 14, 56, 18, 9, 2]
       SCREAMING_SNAKE_CASE_      		=     self.tokenizer.decode(_lowerCAmelCase							,							skip_special_tokens=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     self.tokenizer.decode(generated_ids[1:]							,							skip_special_tokens=_lowerCAmelCase		)
       self.assertEqual(_lowerCAmelCase							,							_lowerCAmelCase		)
       self.assertNotIn(self.tokenizer.eos_token							,							_lowerCAmelCase		)



 def  lowerCAmelCase_    (				self    :						Dict		):
       SCREAMING_SNAKE_CASE_      		=     ['def sum(a,b,c):NEW_LINE_INDENTreturn sum([a,b,c])' * 20]
       self.assertIsInstance(src_text[0]							,							_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     10
       SCREAMING_SNAKE_CASE_      		=     self.tokenizer(_lowerCAmelCase							,							max_length=_lowerCAmelCase							,							truncation=_lowerCAmelCase		).input_ids[0]
       self.assertEqual(ids[-2]							,							2		)
       self.assertEqual(ids[-1]							,							_lowerCAmelCase		)
       self.assertEqual(len(_lowerCAmelCase		)							,							_lowerCAmelCase		)



 def  lowerCAmelCase_    (				self    :						Union[str, Any]		):
       self.assertListEqual(self.tokenizer.convert_tokens_to_ids(['<mask>', '__java__']		)							,							[50_004, 50_001]		)



 def  lowerCAmelCase_    (				self    :						Union[str, Any]		):
       SCREAMING_SNAKE_CASE_      		=     tempfile.mkdtemp()
       SCREAMING_SNAKE_CASE_      		=     self.tokenizer.fairseq_tokens_to_ids
       self.tokenizer.save_pretrained(_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     PLBartTokenizer.from_pretrained(_lowerCAmelCase		)
       self.assertDictEqual(new_tok.fairseq_tokens_to_ids							,							_lowerCAmelCase		)



 @require_torch
 def  lowerCAmelCase_    (				self    :						Dict		):
       SCREAMING_SNAKE_CASE_      		=     self.tokenizer(self.src_text							,							text_target=self.tgt_text							,							padding=_lowerCAmelCase							,							return_tensors='pt'		)
       SCREAMING_SNAKE_CASE_      		=     shift_tokens_right(batch['labels']							,							self.tokenizer.pad_token_id		)

       # fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
       self.assertEqual(batch.input_ids[1][-2:].tolist()							,							[2, PYTHON_CODE]		)
       self.assertEqual(batch.decoder_input_ids[1][0]							,							_lowerCAmelCase		)
       self.assertEqual(batch.decoder_input_ids[1][-1]							,							2		)
       self.assertEqual(batch.labels[1][-2:].tolist()							,							[2, EN_CODE]		)



 @require_torch
 def  lowerCAmelCase_    (				self    :						Any		):
       SCREAMING_SNAKE_CASE_      		=     self.tokenizer(
           self.src_text							,							text_target=self.tgt_text							,							padding=_lowerCAmelCase							,							truncation=_lowerCAmelCase							,							max_length=len(self.expected_src_tokens		)							,							return_tensors='pt'							,							)
       SCREAMING_SNAKE_CASE_      		=     shift_tokens_right(batch['labels']							,							self.tokenizer.pad_token_id		)

       self.assertIsInstance(_lowerCAmelCase							,							_lowerCAmelCase		)

       self.assertEqual((2, 26)							,							batch.input_ids.shape		)
       self.assertEqual((2, 26)							,							batch.attention_mask.shape		)
       SCREAMING_SNAKE_CASE_      		=     batch.input_ids.tolist()[0]
       self.assertListEqual(self.expected_src_tokens							,							_lowerCAmelCase		)
       self.assertEqual(2							,							batch.decoder_input_ids[0, -1]		)  # EOS
       # Test that special tokens are reset
       self.assertEqual(self.tokenizer.prefix_tokens							,							[]		)
       self.assertEqual(self.tokenizer.suffix_tokens							,							[self.tokenizer.eos_token_id, PYTHON_CODE]		)



 def  lowerCAmelCase_    (				self    :						Any		):
       SCREAMING_SNAKE_CASE_      		=     self.tokenizer(self.src_text							,							padding=_lowerCAmelCase							,							truncation=_lowerCAmelCase							,							max_length=3							,							return_tensors='pt'		)
       SCREAMING_SNAKE_CASE_      		=     self.tokenizer(
           text_target=self.tgt_text							,							padding=_lowerCAmelCase							,							truncation=_lowerCAmelCase							,							max_length=10							,							return_tensors='pt'		)
       SCREAMING_SNAKE_CASE_      		=     targets['input_ids']
       SCREAMING_SNAKE_CASE_      		=     shift_tokens_right(_lowerCAmelCase							,							self.tokenizer.pad_token_id		)

       self.assertEqual(batch.input_ids.shape[1]							,							3		)
       self.assertEqual(batch.decoder_input_ids.shape[1]							,							10		)




 @require_torch
 def  lowerCAmelCase_    (				self    :						List[str]		):
       SCREAMING_SNAKE_CASE_      		=     self.tokenizer._build_translation_inputs(
           'A test'							,							return_tensors='pt'							,							src_lang='en_XX'							,							tgt_lang='java'		)

       self.assertEqual(
           nested_simplify(_lowerCAmelCase		)							,							{
               # A, test, EOS, en_XX
               'input_ids': [[150, 242, 2, 50_003]],
               'attention_mask': [[1, 1, 1, 1]],
               # java
               'forced_bos_token_id': 50_001,
           }							,							)


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     (				_a		):






					'''simple docstring'''



					def    _UpperCAmelCase    (				self						):
						UpperCAmelCase__:	Optional[Any]									= pa.array(TypedSequence([1, 2, 3]						)						)
						self.assertEqual(arr.type				,      pa.intaa()						)





					def    _UpperCAmelCase    (				self						):
						with self.assertRaises(lowerCamelCase__						):
							UpperCAmelCase__:	Optional[Any]									= pa.array(TypedSequence([1, 2, 3]						)				,      type=pa.intaa()						)





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





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





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





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





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





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





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





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





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





					@require_pil
					def    _UpperCAmelCase    (				self						):
						import PIL.Image

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



def 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	UpperCAmelCase__:	Optional[Any]									= pa.BufferReader(SCREAMING_SNAKE_CASE      ) if isinstance(SCREAMING_SNAKE_CASE	,pa.Buffer      ) else pa.memory_map(SCREAMING_SNAKE_CASE      )
	UpperCAmelCase__:	str									= pa.ipc.open_stream(SCREAMING_SNAKE_CASE      )
	UpperCAmelCase__:	pa.Table									= 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 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	UpperCAmelCase__:	List[str]									= pa.BufferOutputStream()
	UpperCAmelCase__:	Tuple									= pa.schema(SCREAMING_SNAKE_CASE      ) if fields else None
	with ArrowWriter(stream=SCREAMING_SNAKE_CASE	,schema=SCREAMING_SNAKE_CASE	,writer_batch_size=SCREAMING_SNAKE_CASE      ) as writer:
		writer.write({"col_1": "foo", "col_2": 1}      )
		writer.write({"col_1": "bar", "col_2": 2}      )
		UpperCAmelCase__							,						UpperCAmelCase__:	List[Any]									= writer.finalize()
	assert num_examples == 2
	assert num_bytes > 0
	if not fields:
		UpperCAmelCase__:	Optional[int]									= {"col_1": pa.string(), "col_2": pa.intaa()}
	assert writer._schema == pa.schema(SCREAMING_SNAKE_CASE	,metadata=writer._schema.metadata      )
	_check_output(output.getvalue()	,expected_num_chunks=num_examples if writer_batch_size == 1 else 1      )



def 			_A	(      ):
	UpperCAmelCase__:	Tuple									= pa.BufferOutputStream()
	UpperCAmelCase__:	Any									= Features({"labels": ClassLabel(names=["neg", "pos"]      )}      )
	with ArrowWriter(stream=SCREAMING_SNAKE_CASE	,features=SCREAMING_SNAKE_CASE      ) as writer:
		writer.write({"labels": 0}      )
		writer.write({"labels": 1}      )
		UpperCAmelCase__							,						UpperCAmelCase__:	List[Any]									= 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__:	Dict									= pa.BufferReader(output.getvalue()      )
	UpperCAmelCase__:	Any									= pa.ipc.open_stream(SCREAMING_SNAKE_CASE      )
	UpperCAmelCase__:	pa.Table									= f.read_all()
	UpperCAmelCase__:	Union[str, Any]									= 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(SCREAMING_SNAKE_CASE      )



@pytest.mark.parametrize("writer_batch_size"	,[None, 1, 1_0]      )
def 			_A	(      SCREAMING_SNAKE_CASE      ):
	UpperCAmelCase__:	List[str]									= pa.BufferOutputStream()
	with ArrowWriter(
	    stream=SCREAMING_SNAKE_CASE	,writer_batch_size=SCREAMING_SNAKE_CASE	,hash_salt="split_name"	,check_duplicates=SCREAMING_SNAKE_CASE	,) as writer:
		with pytest.raises(SCREAMING_SNAKE_CASE      ):
			writer.write({"col_1": "foo", "col_2": 1}	,key=[1, 2]      )
			UpperCAmelCase__							,						UpperCAmelCase__:	Any									= writer.finalize()



@pytest.mark.parametrize("writer_batch_size"	,[None, 2, 1_0]      )
def 			_A	(      SCREAMING_SNAKE_CASE      ):
	UpperCAmelCase__:	int									= pa.BufferOutputStream()
	with ArrowWriter(
	    stream=SCREAMING_SNAKE_CASE	,writer_batch_size=SCREAMING_SNAKE_CASE	,hash_salt="split_name"	,check_duplicates=SCREAMING_SNAKE_CASE	,) as writer:
		with pytest.raises(SCREAMING_SNAKE_CASE      ):
			writer.write({"col_1": "foo", "col_2": 1}	,key=1_0      )
			writer.write({"col_1": "bar", "col_2": 2}	,key=1_0      )
			UpperCAmelCase__							,						UpperCAmelCase__:	str									= writer.finalize()



@pytest.mark.parametrize("writer_batch_size"	,[None, 2, 1_0]      )
def 			_A	(      SCREAMING_SNAKE_CASE      ):
	UpperCAmelCase__:	Dict									= pa.BufferOutputStream()
	with ArrowWriter(
	    stream=SCREAMING_SNAKE_CASE	,writer_batch_size=SCREAMING_SNAKE_CASE	,hash_salt="split_name"	,check_duplicates=SCREAMING_SNAKE_CASE	,) as writer:
		writer.write({"col_1": "foo", "col_2": 1}	,key=1      )
		writer.write({"col_1": "bar", "col_2": 2}	,key=2      )
		UpperCAmelCase__							,						UpperCAmelCase__:	Union[str, Any]									= 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 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	UpperCAmelCase__:	int									= pa.BufferOutputStream()
	UpperCAmelCase__:	Optional[Any]									= pa.schema(SCREAMING_SNAKE_CASE      ) if fields else None
	with ArrowWriter(stream=SCREAMING_SNAKE_CASE	,schema=SCREAMING_SNAKE_CASE	,writer_batch_size=SCREAMING_SNAKE_CASE      ) as writer:
		writer.write_batch({"col_1": ["foo", "bar"], "col_2": [1, 2]}      )
		writer.write_batch({"col_1": [], "col_2": []}      )
		UpperCAmelCase__							,						UpperCAmelCase__:	int									= writer.finalize()
	assert num_examples == 2
	assert num_bytes > 0
	if not fields:
		UpperCAmelCase__:	Optional[Any]									= {"col_1": pa.string(), "col_2": pa.intaa()}
	assert writer._schema == pa.schema(SCREAMING_SNAKE_CASE	,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 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	UpperCAmelCase__:	Tuple									= pa.BufferOutputStream()
	UpperCAmelCase__:	List[str]									= pa.schema(SCREAMING_SNAKE_CASE      ) if fields else None
	with ArrowWriter(stream=SCREAMING_SNAKE_CASE	,schema=SCREAMING_SNAKE_CASE	,writer_batch_size=SCREAMING_SNAKE_CASE      ) as writer:
		writer.write_table(pa.Table.from_pydict({"col_1": ["foo", "bar"], "col_2": [1, 2]}      )      )
		UpperCAmelCase__							,						UpperCAmelCase__:	Any									= writer.finalize()
	assert num_examples == 2
	assert num_bytes > 0
	if not fields:
		UpperCAmelCase__:	str									= {"col_1": pa.string(), "col_2": pa.intaa()}
	assert writer._schema == pa.schema(SCREAMING_SNAKE_CASE	,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 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	UpperCAmelCase__:	List[str]									= pa.BufferOutputStream()
	UpperCAmelCase__:	List[Any]									= pa.schema(SCREAMING_SNAKE_CASE      ) if fields else None
	with ArrowWriter(stream=SCREAMING_SNAKE_CASE	,schema=SCREAMING_SNAKE_CASE	,writer_batch_size=SCREAMING_SNAKE_CASE      ) 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__:	Any									= writer.finalize()
	assert num_examples == 2
	assert num_bytes > 0
	if not fields:
		UpperCAmelCase__:	Dict									= {"col_1": pa.string(), "col_2": pa.intaa()}
	assert writer._schema == pa.schema(SCREAMING_SNAKE_CASE	,metadata=writer._schema.metadata      )
	_check_output(output.getvalue()	,expected_num_chunks=num_examples if writer_batch_size == 1 else 1      )



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



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



def 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	if isinstance(lst[0]	,SCREAMING_SNAKE_CASE      ):
		change_first_primitive_element_in_list(lst[0]	,SCREAMING_SNAKE_CASE      )
	else:
		UpperCAmelCase__:	Dict									= 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 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	UpperCAmelCase__:	str									= pa.array(TypedSequence(SCREAMING_SNAKE_CASE	,optimized_int_type=SCREAMING_SNAKE_CASE      )      )
	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 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	# in range
	UpperCAmelCase__:	List[Any]									= pa.array(OptimizedTypedSequence(SCREAMING_SNAKE_CASE	,col=SCREAMING_SNAKE_CASE      )      )
	assert get_base_dtype(arr.type      ) == expected_dtype

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



@pytest.mark.parametrize("raise_exception"	,[False, True]      )
def 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	UpperCAmelCase__:	List[Any]									= str(tmp_path / "dataset-train.arrow"      )
	try:
		with ArrowWriter(path=SCREAMING_SNAKE_CASE      ) as writer:
			if raise_exception:
				raise pa.lib.ArrowInvalid()
			else:
				writer.stream.close()
	except pa.lib.ArrowInvalid:
		pass
	finally:
		assert writer.stream.closed



def 			_A	(      SCREAMING_SNAKE_CASE      ):
	UpperCAmelCase__:	List[Any]									= "mock://dataset-train.arrow"
	with ArrowWriter(path=SCREAMING_SNAKE_CASE	,storage_options=mockfs.storage_options      ) as writer:
		assert isinstance(writer._fs	,type(SCREAMING_SNAKE_CASE      )      )
		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__:	List[Any]									= writer.finalize()
	assert num_examples == 2
	assert num_bytes > 0
	assert mockfs.exists(SCREAMING_SNAKE_CASE      )



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



@require_pil
@pytest.mark.parametrize("embed_local_files"	,[False, True]      )
def 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	import PIL.Image

	UpperCAmelCase__:	List[str]									= str(tmp_path / "test_image_rgb.jpg"      )
	PIL.Image.fromarray(np.zeros((5, 5)	,dtype=np.uinta      )      ).save(SCREAMING_SNAKE_CASE	,format="png"      )
	UpperCAmelCase__:	List[Any]									= pa.BufferOutputStream()
	with ParquetWriter(
	    stream=SCREAMING_SNAKE_CASE	,features=Features({"image": Image()}      )	,embed_local_files=SCREAMING_SNAKE_CASE      ) as writer:
		writer.write({"image": image_path}      )
		writer.finalize()
	UpperCAmelCase__:	Optional[int]									= pa.BufferReader(output.getvalue()      )
	UpperCAmelCase__:	pa.Table									= pq.read_table(SCREAMING_SNAKE_CASE      )
	UpperCAmelCase__:	Dict									= pa_table.to_pydict()
	if embed_local_files:
		assert isinstance(out["image"][0]["path"]	,SCREAMING_SNAKE_CASE      )
		with open(SCREAMING_SNAKE_CASE	,"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 			_A	(      ):
	UpperCAmelCase__:	List[str]									= pa.schema([pa.field("col_1"	,pa.string()	,nullable=SCREAMING_SNAKE_CASE      )]      )
	UpperCAmelCase__:	Optional[Any]									= pa.BufferOutputStream()
	with ArrowWriter(stream=SCREAMING_SNAKE_CASE      ) as writer:
		writer._build_writer(inferred_schema=SCREAMING_SNAKE_CASE      )
		assert writer._schema == pa.schema([pa.field("col_1"	,pa.string()      )]      )