

import argparse
import gc
import json
import os
import shutil
import warnings

import torch

from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer


try:
					from transformers import LlamaTokenizerFast
except ImportError as e:
					warnings.warn(e)
					warnings.warn(
					    """The converted tokenizer will be the `slow` tokenizer. To use the fast, update your `tokenizers` library and re-run the tokenizer conversion"""
					)
					_lowerCAmelCase  :      Tuple			 =None
_lowerCAmelCase  :      Any			 ={
    """7B""": 1_10_08,
    """13B""": 1_38_24,
    """30B""": 1_79_20,
    """65B""": 2_20_16,
    """70B""": 2_86_72,
}
_lowerCAmelCase  :      Optional[Any]			 ={
    """7B""": 1,
    """7Bf""": 1,
    """13B""": 2,
    """13Bf""": 2,
    """30B""": 4,
    """65B""": 8,
    """70B""": 8,
    """70Bf""": 8,
}



def 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE=1	,SCREAMING_SNAKE_CASE=2_5_6      ):
	return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3      )      ) + multiple_of - 1) // multiple_of)



def 			_A	(      SCREAMING_SNAKE_CASE      ):
	with open(SCREAMING_SNAKE_CASE	,"r"      ) as f:
		return json.load(SCREAMING_SNAKE_CASE      )



def 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	with open(SCREAMING_SNAKE_CASE	,"w"      ) as f:
		json.dump(SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      )



def 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE=True      ):
	os.makedirs(SCREAMING_SNAKE_CASE	,exist_ok=SCREAMING_SNAKE_CASE      )
	UpperCAmelCase__:	List[str]									= os.path.join(SCREAMING_SNAKE_CASE	,"tmp"      )
	os.makedirs(SCREAMING_SNAKE_CASE	,exist_ok=SCREAMING_SNAKE_CASE      )

	UpperCAmelCase__:	str									= read_json(os.path.join(SCREAMING_SNAKE_CASE	,"params.json"      )      )
	UpperCAmelCase__:	Optional[int]									= NUM_SHARDS[model_size]
	UpperCAmelCase__:	Tuple									= params["n_layers"]
	UpperCAmelCase__:	Tuple									= params["n_heads"]
	UpperCAmelCase__:	Optional[Any]									= n_heads // num_shards
	UpperCAmelCase__:	str									= params["dim"]
	UpperCAmelCase__:	int									= dim // n_heads
	UpperCAmelCase__:	List[str]									= 1_00_00.0
	UpperCAmelCase__:	List[str]									= 1.0 / (base ** (torch.arange(0	,SCREAMING_SNAKE_CASE	,2      ).float() / dims_per_head))

	if "n_kv_heads" in params:
		UpperCAmelCase__:	Dict									= params["n_kv_heads"]  # for GQA / MQA
		UpperCAmelCase__:	int									= n_heads_per_shard // num_key_value_heads
		UpperCAmelCase__:	Union[str, Any]									= dim // num_key_value_heads
	else:  # compatibility with other checkpoints
		UpperCAmelCase__:	Dict									= n_heads
		UpperCAmelCase__:	Optional[int]									= n_heads_per_shard
		UpperCAmelCase__:	Tuple									= dim

	# permute for sliced rotary
	def permute(SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE=n_heads	,SCREAMING_SNAKE_CASE=dim	,SCREAMING_SNAKE_CASE=dim      ):
		return w.view(SCREAMING_SNAKE_CASE	,dima // n_heads // 2	,2	,SCREAMING_SNAKE_CASE      ).transpose(1	,2      ).reshape(SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      )

	print(f"Fetching all parameters from the checkpoint at {input_base_path}."      )
	# Load weights
	if model_size == "7B":
		# Not sharded
		# (The sharded implementation would also work, but this is simpler.)
		UpperCAmelCase__:	List[Any]									= torch.load(os.path.join(SCREAMING_SNAKE_CASE	,"consolidated.00.pth"      )	,map_location="cpu"      )
	else:
		# Sharded
		UpperCAmelCase__:	Any									= [
		    torch.load(os.path.join(SCREAMING_SNAKE_CASE	,f"consolidated.{i:02d}.pth"      )	,map_location="cpu"      )
		    for i in range(SCREAMING_SNAKE_CASE      )
		]
	UpperCAmelCase__:	int									= 0
	UpperCAmelCase__:	Dict									= {"weight_map": {}}
	for layer_i in range(SCREAMING_SNAKE_CASE      ):
		UpperCAmelCase__:	Union[str, Any]									= f"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin"
		if model_size == "7B":
			# Unsharded
			UpperCAmelCase__:	int									= {
			    f"model.layers.{layer_i}.self_attn.q_proj.weight": permute(
			        loaded[f"layers.{layer_i}.attention.wq.weight"]      ),
			    f"model.layers.{layer_i}.self_attn.k_proj.weight": permute(
			        loaded[f"layers.{layer_i}.attention.wk.weight"]      ),
			    f"model.layers.{layer_i}.self_attn.v_proj.weight": loaded[f"layers.{layer_i}.attention.wv.weight"],
			    f"model.layers.{layer_i}.self_attn.o_proj.weight": loaded[f"layers.{layer_i}.attention.wo.weight"],
			    f"model.layers.{layer_i}.mlp.gate_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w1.weight"],
			    f"model.layers.{layer_i}.mlp.down_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w2.weight"],
			    f"model.layers.{layer_i}.mlp.up_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w3.weight"],
			    f"model.layers.{layer_i}.input_layernorm.weight": loaded[f"layers.{layer_i}.attention_norm.weight"],
			    f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[f"layers.{layer_i}.ffn_norm.weight"],
			}
		else:
			# Sharded
			# Note that attention.w{q,k,v,o}, feed_fordward.w[1,2,3], attention_norm.weight and ffn_norm.weight share
			# the same storage object, saving attention_norm and ffn_norm will save other weights too, which is
			# redundant as other weights will be stitched from multiple shards. To avoid that, they are cloned.

			UpperCAmelCase__:	int									= {
			    f"model.layers.{layer_i}.input_layernorm.weight": loaded[0][
			        f"layers.{layer_i}.attention_norm.weight"
			    ].clone(),
			    f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[0][
			        f"layers.{layer_i}.ffn_norm.weight"
			    ].clone(),
			}
			UpperCAmelCase__:	Dict									= permute(
			    torch.cat(
			        [
			            loaded[i][f"layers.{layer_i}.attention.wq.weight"].view(SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      )
			            for i in range(SCREAMING_SNAKE_CASE      )
			        ]	,dim=0	,).reshape(SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      )      )
			UpperCAmelCase__:	List[Any]									= permute(
			    torch.cat(
			        [
			            loaded[i][f"layers.{layer_i}.attention.wk.weight"].view(
			                SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      )
			            for i in range(SCREAMING_SNAKE_CASE      )
			        ]	,dim=0	,).reshape(SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      )	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,)
			UpperCAmelCase__:	Union[str, Any]									= torch.cat(
			    [
			        loaded[i][f"layers.{layer_i}.attention.wv.weight"].view(
			            SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      )
			        for i in range(SCREAMING_SNAKE_CASE      )
			    ]	,dim=0	,).reshape(SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      )

			UpperCAmelCase__:	Tuple									= torch.cat(
			    [loaded[i][f"layers.{layer_i}.attention.wo.weight"] for i in range(SCREAMING_SNAKE_CASE      )]	,dim=1      )
			UpperCAmelCase__:	Dict									= torch.cat(
			    [loaded[i][f"layers.{layer_i}.feed_forward.w1.weight"] for i in range(SCREAMING_SNAKE_CASE      )]	,dim=0      )
			UpperCAmelCase__:	Optional[int]									= torch.cat(
			    [loaded[i][f"layers.{layer_i}.feed_forward.w2.weight"] for i in range(SCREAMING_SNAKE_CASE      )]	,dim=1      )
			UpperCAmelCase__:	Dict									= torch.cat(
			    [loaded[i][f"layers.{layer_i}.feed_forward.w3.weight"] for i in range(SCREAMING_SNAKE_CASE      )]	,dim=0      )

		UpperCAmelCase__:	Union[str, Any]									= inv_freq
		for k, v in state_dict.items():
			UpperCAmelCase__:	str									= filename
			param_count += v.numel()
		torch.save(SCREAMING_SNAKE_CASE	,os.path.join(SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      )      )

	UpperCAmelCase__:	Optional[Any]									= f"pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin"
	if model_size == "7B":
		# Unsharded
		UpperCAmelCase__:	str									= {
		    "model.embed_tokens.weight": loaded["tok_embeddings.weight"],
		    "model.norm.weight": loaded["norm.weight"],
		    "lm_head.weight": loaded["output.weight"],
		}
	else:
		UpperCAmelCase__:	Optional[int]									= {
		    "model.norm.weight": loaded[0]["norm.weight"],
		    "model.embed_tokens.weight": torch.cat(
		        [loaded[i]["tok_embeddings.weight"] for i in range(SCREAMING_SNAKE_CASE      )]	,dim=1      ),
		    "lm_head.weight": torch.cat([loaded[i]["output.weight"] for i in range(SCREAMING_SNAKE_CASE      )]	,dim=0      ),
		}

	for k, v in state_dict.items():
		UpperCAmelCase__:	int									= filename
		param_count += v.numel()
	torch.save(SCREAMING_SNAKE_CASE	,os.path.join(SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      )      )

	# Write configs
	UpperCAmelCase__:	Tuple									= {"total_size": param_count * 2}
	write_json(SCREAMING_SNAKE_CASE	,os.path.join(SCREAMING_SNAKE_CASE	,"pytorch_model.bin.index.json"      )      )
	UpperCAmelCase__:	int									= params["ffn_dim_multiplier"] if "ffn_dim_multiplier" in params else 1
	UpperCAmelCase__:	str									= params["multiple_of"] if "multiple_of" in params else 2_5_6
	UpperCAmelCase__:	Dict									= LlamaConfig(
	    hidden_size=SCREAMING_SNAKE_CASE	,intermediate_size=compute_intermediate_size(SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      )	,num_attention_heads=params["n_heads"]	,num_hidden_layers=params["n_layers"]	,rms_norm_eps=params["norm_eps"]	,num_key_value_heads=SCREAMING_SNAKE_CASE	,)
	config.save_pretrained(SCREAMING_SNAKE_CASE      )

	# Make space so we can load the model properly now.
	del state_dict
	del loaded
	gc.collect()

	print("Loading the checkpoint in a Llama model."      )
	UpperCAmelCase__:	Any									= LlamaForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE	,torch_dtype=torch.floataa	,low_cpu_mem_usage=SCREAMING_SNAKE_CASE      )
	# Avoid saving this as part of the config.
	del model.config._name_or_path

	print("Saving in the Transformers format."      )
	model.save_pretrained(SCREAMING_SNAKE_CASE	,safe_serialization=SCREAMING_SNAKE_CASE      )
	shutil.rmtree(SCREAMING_SNAKE_CASE      )



def 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	# Initialize the tokenizer based on the `spm` model
	UpperCAmelCase__:	Tuple									= LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
	print(f"Saving a {tokenizer_class.__name__} to {tokenizer_path}."      )
	UpperCAmelCase__:	str									= tokenizer_class(SCREAMING_SNAKE_CASE      )
	tokenizer.save_pretrained(SCREAMING_SNAKE_CASE      )



def 			_A	(      ):
	UpperCAmelCase__:	str									= argparse.ArgumentParser()
	parser.add_argument(
	    "--input_dir"	,help="Location of LLaMA weights, which contains tokenizer.model and model folders"	,)
	parser.add_argument(
	    "--model_size"	,choices=["7B", "7Bf", "13B", "13Bf", "30B", "65B", "70B", "70Bf", "tokenizer_only"]	,)
	parser.add_argument(
	    "--output_dir"	,help="Location to write HF model and tokenizer"	,)
	parser.add_argument("--safe_serialization"	,type=SCREAMING_SNAKE_CASE	,help="Whether or not to save using `safetensors`."      )
	UpperCAmelCase__:	str									= parser.parse_args()
	if args.model_size != "tokenizer_only":
		write_model(
		    model_path=args.output_dir	,input_base_path=os.path.join(args.input_dir	,args.model_size      )	,model_size=args.model_size	,safe_serialization=args.safe_serialization	,)
	UpperCAmelCase__:	Optional[Any]									= os.path.join(args.input_dir	,"tokenizer.model"      )
	write_tokenizer(args.output_dir	,SCREAMING_SNAKE_CASE      )


if __name__ == "__main__":
					main()


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()      )]      )