
import unittest

from transformers import ElectraConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
	import torch

	from transformers import (
	    MODEL_FOR_PRETRAINING_MAPPING,
	    ElectraForCausalLM,
	    ElectraForMaskedLM,
	    ElectraForMultipleChoice,
	    ElectraForPreTraining,
	    ElectraForQuestionAnswering,
	    ElectraForSequenceClassification,
	    ElectraForTokenClassification,
	    ElectraModel,
	)
	from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST





class 		__lowerCamelCase     :



							def __init__(     self	,				__snake_case	,				__snake_case=1_3	,				__snake_case=7	,				__snake_case=True	,				__snake_case=True	,				__snake_case=True	,				__snake_case=True	,				__snake_case=9_9	,				__snake_case=3_2	,				__snake_case=5	,				__snake_case=4	,				__snake_case=3_7	,				__snake_case="gelu"	,				__snake_case=0.1	,				__snake_case=0.1	,				__snake_case=5_1_2	,				__snake_case=1_6	,				__snake_case=2	,				__snake_case=0.02	,				__snake_case=3	,				__snake_case=4	,				__snake_case=None	,				)						->     Any:






								"""simple docstring"""



								UpperCAmelCase:					str 						=     parent
								UpperCAmelCase:					Dict 						=     batch_size
								UpperCAmelCase:					Union[str, Any] 						=     seq_length
								UpperCAmelCase:					Tuple 						=     is_training
								UpperCAmelCase:					List[Any] 						=     use_input_mask
								UpperCAmelCase:					str 						=     use_token_type_ids
								UpperCAmelCase:					Union[str, Any] 						=     use_labels
								UpperCAmelCase:					List[Any] 						=     vocab_size
								UpperCAmelCase:					Any 						=     hidden_size
								UpperCAmelCase:					int 						=     num_hidden_layers
								UpperCAmelCase:					str 						=     num_attention_heads
								UpperCAmelCase:					Optional[int] 						=     intermediate_size
								UpperCAmelCase:					Tuple 						=     hidden_act
								UpperCAmelCase:					Any 						=     hidden_dropout_prob
								UpperCAmelCase:					Optional[int] 						=     attention_probs_dropout_prob
								UpperCAmelCase:					Dict 						=     max_position_embeddings
								UpperCAmelCase:					Optional[int] 						=     type_vocab_size
								UpperCAmelCase:					Optional[Any] 						=     type_sequence_label_size
								UpperCAmelCase:					Union[str, Any] 						=     initializer_range
								UpperCAmelCase:					int 						=     num_labels
								UpperCAmelCase:					List[Any] 						=     num_choices
								UpperCAmelCase:					Dict 						=     scope



							def     A__ (     self							)						->     Optional[int]:






								"""simple docstring"""



								UpperCAmelCase:					str 						=     ids_tensor([self.batch_size, self.seq_length]	,				self.vocab_size							)

								UpperCAmelCase:					Union[str, Any] 						=     None
								if self.use_input_mask:
									UpperCAmelCase:					str 						=     random_attention_mask([self.batch_size, self.seq_length]							)

								UpperCAmelCase:					int 						=     None
								if self.use_token_type_ids:
									UpperCAmelCase:					Union[str, Any] 						=     ids_tensor([self.batch_size, self.seq_length]	,				self.type_vocab_size							)

								UpperCAmelCase:					str 						=     None
								UpperCAmelCase:					Any 						=     None
								UpperCAmelCase:					List[Any] 						=     None
								if self.use_labels:
									UpperCAmelCase:					str 						=     ids_tensor([self.batch_size]	,				self.type_sequence_label_size							)
									UpperCAmelCase:					List[str] 						=     ids_tensor([self.batch_size, self.seq_length]	,				self.num_labels							)
									UpperCAmelCase:					Dict 						=     ids_tensor([self.batch_size]	,				self.num_choices							)
									UpperCAmelCase:					Tuple 						=     ids_tensor([self.batch_size, self.seq_length]	,				1							)

								UpperCAmelCase:					List[Any] 						=     self.get_config()

								return (
								    config,
								    input_ids,
								    token_type_ids,
								    input_mask,
								    sequence_labels,
								    token_labels,
								    choice_labels,
								    fake_token_labels,
								)



							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								return ElectraConfig(
								    vocab_size=self.vocab_size	,				hidden_size=self.hidden_size	,				num_hidden_layers=self.num_hidden_layers	,				num_attention_heads=self.num_attention_heads	,				intermediate_size=self.intermediate_size	,				hidden_act=self.hidden_act	,				hidden_dropout_prob=self.hidden_dropout_prob	,				attention_probs_dropout_prob=self.attention_probs_dropout_prob	,				max_position_embeddings=self.max_position_embeddings	,				type_vocab_size=self.type_vocab_size	,				is_decoder=__snake_case	,				initializer_range=self.initializer_range	,				)



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






								"""simple docstring"""



								(
								    (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   
								):					str 						=     self.prepare_config_and_inputs()

								UpperCAmelCase:					Union[str, Any] 						=     True
								UpperCAmelCase:					Any 						=     floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
								UpperCAmelCase:					int 						=     ids_tensor([self.batch_size, self.seq_length]	,				vocab_size=2							)

								return (
								    config,
								    input_ids,
								    token_type_ids,
								    input_mask,
								    sequence_labels,
								    token_labels,
								    choice_labels,
								    encoder_hidden_states,
								    encoder_attention_mask,
								)



							def     A__ (     self	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				)						->     List[Any]:






								"""simple docstring"""



								UpperCAmelCase:					Dict 						=     ElectraModel(config=__snake_case							)
								model.to(__snake_case							)
								model.eval()
								UpperCAmelCase:					Any 						=     model(__snake_case	,				attention_mask=__snake_case	,				token_type_ids=__snake_case							)
								UpperCAmelCase:					Optional[Any] 						=     model(__snake_case	,				token_type_ids=__snake_case							)
								UpperCAmelCase:					List[Any] 						=     model(__snake_case							)
								self.parent.assertEqual(result.last_hidden_state.shape	,				(self.batch_size, self.seq_length, self.hidden_size)							)



							def     A__ (     self	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				)						->     str:






								"""simple docstring"""



								UpperCAmelCase:					Union[str, Any] 						=     True
								UpperCAmelCase:					str 						=     ElectraModel(__snake_case							)
								model.to(__snake_case							)
								model.eval()
								UpperCAmelCase:					Optional[int] 						=     model(
								    __snake_case	,				attention_mask=__snake_case	,				token_type_ids=__snake_case	,				encoder_hidden_states=__snake_case	,				encoder_attention_mask=__snake_case	,				)
								UpperCAmelCase:					Tuple 						=     model(
								    __snake_case	,				attention_mask=__snake_case	,				token_type_ids=__snake_case	,				encoder_hidden_states=__snake_case	,				)
								UpperCAmelCase:					int 						=     model(__snake_case	,				attention_mask=__snake_case	,				token_type_ids=__snake_case							)
								self.parent.assertEqual(result.last_hidden_state.shape	,				(self.batch_size, self.seq_length, self.hidden_size)							)



							def     A__ (     self	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				)						->     List[Any]:






								"""simple docstring"""



								UpperCAmelCase:					Optional[Any] 						=     ElectraForMaskedLM(config=__snake_case							)
								model.to(__snake_case							)
								model.eval()
								UpperCAmelCase:					List[Any] 						=     model(__snake_case	,				attention_mask=__snake_case	,				token_type_ids=__snake_case	,				labels=__snake_case							)
								self.parent.assertEqual(result.logits.shape	,				(self.batch_size, self.seq_length, self.vocab_size)							)



							def     A__ (     self	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				)						->     Dict:






								"""simple docstring"""



								UpperCAmelCase:					int 						=     ElectraForCausalLM(config=__snake_case							)
								model.to(__snake_case							)
								model.eval()
								UpperCAmelCase:					List[Any] 						=     model(__snake_case	,				attention_mask=__snake_case	,				token_type_ids=__snake_case	,				labels=__snake_case							)
								self.parent.assertEqual(result.logits.shape	,				(self.batch_size, self.seq_length, self.vocab_size)							)



							def     A__ (     self	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				)						->     List[str]:






								"""simple docstring"""



								UpperCAmelCase:					str 						=     self.num_labels
								UpperCAmelCase:					Dict 						=     ElectraForTokenClassification(config=__snake_case							)
								model.to(__snake_case							)
								model.eval()
								UpperCAmelCase:					List[str] 						=     model(__snake_case	,				attention_mask=__snake_case	,				token_type_ids=__snake_case	,				labels=__snake_case							)
								self.parent.assertEqual(result.logits.shape	,				(self.batch_size, self.seq_length, self.num_labels)							)



							def     A__ (     self	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				)						->     Dict:






								"""simple docstring"""



								UpperCAmelCase:					Union[str, Any] 						=     self.num_labels
								UpperCAmelCase:					Optional[int] 						=     ElectraForPreTraining(config=__snake_case							)
								model.to(__snake_case							)
								model.eval()
								UpperCAmelCase:					Dict 						=     model(__snake_case	,				attention_mask=__snake_case	,				token_type_ids=__snake_case	,				labels=__snake_case							)
								self.parent.assertEqual(result.logits.shape	,				(self.batch_size, self.seq_length)							)



							def     A__ (     self	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				)						->     Any:






								"""simple docstring"""



								UpperCAmelCase:					int 						=     self.num_labels
								UpperCAmelCase:					Union[str, Any] 						=     ElectraForSequenceClassification(__snake_case							)
								model.to(__snake_case							)
								model.eval()
								UpperCAmelCase:					Optional[int] 						=     model(__snake_case	,				attention_mask=__snake_case	,				token_type_ids=__snake_case	,				labels=__snake_case							)
								self.parent.assertEqual(result.logits.shape	,				(self.batch_size, self.num_labels)							)



							def     A__ (     self	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				)						->     str:






								"""simple docstring"""



								UpperCAmelCase:					Dict 						=     ElectraForQuestionAnswering(config=__snake_case							)
								model.to(__snake_case							)
								model.eval()
								UpperCAmelCase:					Optional[Any] 						=     model(
								    __snake_case	,				attention_mask=__snake_case	,				token_type_ids=__snake_case	,				start_positions=__snake_case	,				end_positions=__snake_case	,				)
								self.parent.assertEqual(result.start_logits.shape	,				(self.batch_size, self.seq_length)							)
								self.parent.assertEqual(result.end_logits.shape	,				(self.batch_size, self.seq_length)							)



							def     A__ (     self	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				__snake_case	,				)						->     Optional[Any]:






								"""simple docstring"""



								UpperCAmelCase:					Optional[Any] 						=     self.num_choices
								UpperCAmelCase:					str 						=     ElectraForMultipleChoice(config=__snake_case							)
								model.to(__snake_case							)
								model.eval()
								UpperCAmelCase:					int 						=     input_ids.unsqueeze(1							).expand(-1	,				self.num_choices	,				-1							).contiguous()
								UpperCAmelCase:					str 						=     token_type_ids.unsqueeze(1							).expand(-1	,				self.num_choices	,				-1							).contiguous()
								UpperCAmelCase:					Any 						=     input_mask.unsqueeze(1							).expand(-1	,				self.num_choices	,				-1							).contiguous()
								UpperCAmelCase:					Union[str, Any] 						=     model(
								    __snake_case	,				attention_mask=__snake_case	,				token_type_ids=__snake_case	,				labels=__snake_case	,				)
								self.parent.assertEqual(result.logits.shape	,				(self.batch_size, self.num_choices)							)




							def     A__ (     self							)						->     Optional[int]:






								"""simple docstring"""



								UpperCAmelCase:					Union[str, Any] 						=     self.prepare_config_and_inputs()
								(
								    (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   (
								    UpperCAmelCase
								)						,   
								):					str 						=     config_and_inputs
								UpperCAmelCase:					int 						=     {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
								return config, inputs_dict





@require_torch
class 		__lowerCamelCase     (   lowercase							, lowercase							, unittest.TestCase		):
							lowerCamelCase__:				str      =							(
							    (
							        ElectraModel,
							        ElectraForPreTraining,
							        ElectraForMaskedLM,
							        ElectraForCausalLM,
							        ElectraForMultipleChoice,
							        ElectraForTokenClassification,
							        ElectraForSequenceClassification,
							        ElectraForQuestionAnswering,
							    )
							    if is_torch_available()
							    else ()
							)
							lowerCamelCase__:				int      =							(
							    {
							        '''feature-extraction''': ElectraModel,
							        '''fill-mask''': ElectraForMaskedLM,
							        '''question-answering''': ElectraForQuestionAnswering,
							        '''text-classification''': ElectraForSequenceClassification,
							        '''text-generation''': ElectraForCausalLM,
							        '''token-classification''': ElectraForTokenClassification,
							        '''zero-shot''': ElectraForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							lowerCamelCase__:				Optional[Any]      =							True



							def     A__ (     self	,				__snake_case	,				__snake_case	,				__snake_case=False							)						->     Any:






								"""simple docstring"""



								UpperCAmelCase:					List[Any] 						=     super()._prepare_for_class(__snake_case	,				__snake_case	,				return_labels=__snake_case							)

								if return_labels:
									if model_class in get_values(__snake_case							):
										UpperCAmelCase:					Union[str, Any] 						=     torch.zeros(
										    (self.model_tester.batch_size, self.model_tester.seq_length)	,				dtype=torch.long	,				device=__snake_case							)
								return inputs_dict



							def     A__ (     self							)						->     Dict:






								"""simple docstring"""



								UpperCAmelCase:					str 						=     ElectraModelTester(self							)
								UpperCAmelCase:					Optional[int] 						=     ConfigTester(self	,				config_class=__snake_case	,				hidden_size=3_7							)



							def     A__ (     self							)						->     int:






								"""simple docstring"""



								self.config_tester.run_common_tests()



							def     A__ (     self							)						->     str:






								"""simple docstring"""



								UpperCAmelCase:					List[Any] 						=     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_electra_model(*__snake_case							)



							def     A__ (     self							)						->     Optional[Any]:






								"""simple docstring"""



								UpperCAmelCase:					Optional[int] 						=     self.model_tester.prepare_config_and_inputs_for_decoder()
								self.model_tester.create_and_check_electra_model_as_decoder(*__snake_case							)



							def     A__ (     self							)						->     Dict:






								"""simple docstring"""



								UpperCAmelCase:					Optional[Any] 						=     self.model_tester.prepare_config_and_inputs()
								for type in ["absolute", "relative_key", "relative_key_query"]:
									UpperCAmelCase:					Dict 						=     type
									self.model_tester.create_and_check_electra_model(*__snake_case							)



							def     A__ (     self							)						->     Optional[Any]:






								"""simple docstring"""



								UpperCAmelCase:					Optional[Any] 						=     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_electra_for_masked_lm(*__snake_case							)



							def     A__ (     self							)						->     str:






								"""simple docstring"""



								UpperCAmelCase:					Union[str, Any] 						=     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_electra_for_token_classification(*__snake_case							)



							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase:					Dict 						=     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_electra_for_pretraining(*__snake_case							)



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






								"""simple docstring"""



								UpperCAmelCase:					Tuple 						=     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_electra_for_sequence_classification(*__snake_case							)



							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase:					str 						=     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_electra_for_question_answering(*__snake_case							)



							def     A__ (     self							)						->     Dict:






								"""simple docstring"""



								UpperCAmelCase:					Tuple 						=     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_electra_for_multiple_choice(*__snake_case							)



							@slow
							def     A__ (     self							)						->     Optional[int]:






								"""simple docstring"""



								for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									UpperCAmelCase:					Any 						=     ElectraModel.from_pretrained(__snake_case							)
									self.assertIsNotNone(__snake_case							)




							def     A__ (     self							)						->     Any:






								"""simple docstring"""



								UpperCAmelCase:					List[str] 						=     self.model_tester.prepare_config_and_inputs_for_decoder()
								self.model_tester.create_and_check_electra_for_causal_lm(*__snake_case							)





@require_torch
class 		__lowerCamelCase     (   unittest.TestCase		):



							@slow
							def     A__ (     self							)						->     int:






								"""simple docstring"""



								UpperCAmelCase:					Union[str, Any] 						=     ElectraModel.from_pretrained("google/electra-small-discriminator"							)
								UpperCAmelCase:					int 						=     torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]							)
								UpperCAmelCase:					List[Any] 						=     torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]							)
								UpperCAmelCase:					int 						=     model(__snake_case	,				attention_mask=__snake_case							)[0]
								UpperCAmelCase:					Union[str, Any] 						=     torch.Size((1, 1_1, 2_5_6)							)
								self.assertEqual(output.shape	,				__snake_case							)
								UpperCAmelCase:					List[Any] 						=     torch.tensor(
								    [[[0.44_71, 0.68_21, -0.32_65], [0.46_27, 0.52_55, -0.36_68], [0.45_32, 0.33_13, -0.43_44]]]							)

								self.assertTrue(torch.allclose(output[:, 1:4, 1:4]	,				__snake_case	,				atol=1e-4							)							)


import os
import tempfile
from pathlib import Path
from unittest import TestCase

import pyarrow as pa
import pytest

from datasets.arrow_dataset import Dataset
from datasets.arrow_reader import ArrowReader, BaseReader, FileInstructions, ReadInstruction, make_file_instructions
from datasets.info import DatasetInfo
from datasets.splits import NamedSplit, Split, SplitDict, SplitInfo

from .utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases





class 		__lowerCamelCase     (   lowercase		):



							def     A__ (     self	,				__snake_case	,				__snake_case=False							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase						,   UpperCAmelCase						,   UpperCAmelCase:					Dict 						=     (
								    filename_skip_take["filename"],
								    filename_skip_take["skip"] if "skip" in filename_skip_take else None,
								    filename_skip_take["take"] if "take" in filename_skip_take else None,
								)
								open(os.path.join(__snake_case							)	,				"wb"							).close()
								UpperCAmelCase:					Any 						=     pa.Table.from_pydict({"filename": [Path(__snake_case							).name] * 1_0_0}							)
								if take == -1:
									UpperCAmelCase:					List[Any] 						=     len(__snake_case							) - skip
								if skip is not None and take is not None:
									UpperCAmelCase:					Optional[int] 						=     pa_table.slice(__snake_case	,				__snake_case							)
								return pa_table





class 		__lowerCamelCase     (   lowercase		):



							def     A__ (     self							)						->     int:






								"""simple docstring"""



								UpperCAmelCase:					Tuple 						=     "my_name"
								UpperCAmelCase:					str 						=     SplitInfo(name="train"	,				num_examples=1_0_0							)
								UpperCAmelCase:					Optional[int] 						=     SplitInfo(name="test"	,				num_examples=1_0_0							)
								UpperCAmelCase:					Optional[int] 						=     [train_info, test_info]
								UpperCAmelCase:					Tuple 						=     SplitDict()
								split_dict.add(__snake_case							)
								split_dict.add(__snake_case							)
								UpperCAmelCase:					Optional[int] 						=     DatasetInfo(splits=__snake_case							)

								with tempfile.TemporaryDirectory() as tmp_dir:
									UpperCAmelCase:					int 						=     ReaderTest(__snake_case	,				__snake_case							)

									UpperCAmelCase:					int 						=     "test[:33%]"
									UpperCAmelCase:					Dict 						=     Dataset(**reader.read(__snake_case	,				__snake_case	,				__snake_case							)							)
									self.assertEqual(dset["filename"][0]	,				F'{name}-test'							)
									self.assertEqual(dset.num_rows	,				3_3							)
									self.assertEqual(dset.num_columns	,				1							)

									UpperCAmelCase:					int 						=     ["train", "test[:33%]"]
									UpperCAmelCase:					str 						=     [Split.TRAIN, ReadInstruction.from_spec("test[:33%]"							)]
									for instructions in [instructionsa, instructionsa]:
										UpperCAmelCase:					List[str] 						=     [reader.read(__snake_case	,				__snake_case	,				__snake_case							) for instr in instructions]
										UpperCAmelCase						,   UpperCAmelCase:					List[Any] 						=     (Dataset(**__snake_case							) for dataset_kwargs in datasets_kwargs)
										self.assertEqual(train_dset["filename"][0]	,				F'{name}-train'							)
										self.assertEqual(train_dset.num_rows	,				1_0_0							)
										self.assertEqual(train_dset.num_columns	,				1							)
										self.assertIsInstance(train_dset.split	,				__snake_case							)
										self.assertEqual(str(train_dset.split							)	,				"train"							)
										self.assertEqual(test_dset["filename"][0]	,				F'{name}-test'							)
										self.assertEqual(test_dset.num_rows	,				3_3							)
										self.assertEqual(test_dset.num_columns	,				1							)
										self.assertIsInstance(test_dset.split	,				__snake_case							)
										self.assertEqual(str(test_dset.split							)	,				"test[:33%]"							)
										del train_dset, test_dset



							def     A__ (     self							)						->     Optional[Any]:






								"""simple docstring"""



								UpperCAmelCase:					str 						=     "my_name"
								UpperCAmelCase:					List[Any] 						=     SplitInfo(name="train"	,				num_examples=1_0_0_0	,				shard_lengths=[1_0_0] * 1_0							)
								UpperCAmelCase:					List[Any] 						=     [train_info]
								UpperCAmelCase:					Dict 						=     SplitDict()
								split_dict.add(__snake_case							)
								UpperCAmelCase:					int 						=     DatasetInfo(splits=__snake_case							)

								with tempfile.TemporaryDirectory() as tmp_dir:
									UpperCAmelCase:					Optional[Any] 						=     ReaderTest(__snake_case	,				__snake_case							)

									UpperCAmelCase:					Optional[Any] 						=     "train[:33%]"
									UpperCAmelCase:					Tuple 						=     Dataset(**reader.read(__snake_case	,				__snake_case	,				__snake_case							)							)
									self.assertEqual(dset["filename"][0]	,				F'{name}-train-00000-of-00010'							)
									self.assertEqual(dset["filename"][-1]	,				F'{name}-train-00003-of-00010'							)
									self.assertEqual(dset.num_rows	,				3_3_0							)
									self.assertEqual(dset.num_columns	,				1							)




							def     A__ (     self							)						->     Any:






								"""simple docstring"""



								UpperCAmelCase:					Union[str, Any] 						=     SplitInfo(name="train"	,				num_examples=1_0_0							)
								UpperCAmelCase:					int 						=     SplitInfo(name="test"	,				num_examples=1_0_0							)
								UpperCAmelCase:					List[Any] 						=     SplitDict()
								split_dict.add(__snake_case							)
								split_dict.add(__snake_case							)
								UpperCAmelCase:					Optional[Any] 						=     DatasetInfo(splits=__snake_case							)

								with tempfile.TemporaryDirectory() as tmp_dir:
									UpperCAmelCase:					List[Any] 						=     ReaderTest(__snake_case	,				__snake_case							)

									UpperCAmelCase:					List[str] 						=     [
									    {"filename": os.path.join(__snake_case	,				"train"							)},
									    {"filename": os.path.join(__snake_case	,				"test"							), "skip": 1_0, "take": 1_0},
									]
									UpperCAmelCase:					Tuple 						=     Dataset(**reader.read_files(__snake_case	,				original_instructions="train+test[10:20]"							)							)
									self.assertEqual(dset.num_rows	,				1_1_0							)
									self.assertEqual(dset.num_columns	,				1							)
									del dset



@pytest.mark.parametrize("in_memory"				,					[False, True]				)
def      __UpperCAmelCase    (  snake_case_							:		str				,					snake_case_							:		str				,					snake_case_							:		Any				):



	'''simple docstring'''



	UpperCAmelCase:					Union[str, Any] 						=     arrow_file
	with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
		UpperCAmelCase:					int 						=     ArrowReader.read_table(snake_case_				,					in_memory=snake_case_				)
	assert table.shape == dataset.data.shape
	assert set(table.column_names				) == set(dataset.data.column_names				)
	assert dict(table.to_pydict()				) == dict(dataset.data.to_pydict()				)  # to_pydict returns OrderedDict



@pytest.mark.parametrize("in_memory"				,					[False, True]				)
def      __UpperCAmelCase    (  snake_case_							:		List[str]				,					snake_case_							:		Optional[int]				,					snake_case_							:		int				):



	'''simple docstring'''



	UpperCAmelCase:					List[Any] 						=     arrow_file
	UpperCAmelCase:					Tuple 						=     ArrowReader(""				,					snake_case_				)
	with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
		UpperCAmelCase:					Any 						=     reader.read_files([{"filename": filename}]				,					in_memory=snake_case_				)
	assert dataset_kwargs.keys() == {"arrow_table", "info", "split"}
	UpperCAmelCase:					Any 						=     dataset_kwargs["arrow_table"]
	assert table.shape == dataset.data.shape
	assert set(table.column_names				) == set(dataset.data.column_names				)
	assert dict(table.to_pydict()				) == dict(dataset.data.to_pydict()				)  # to_pydict returns OrderedDict



def      __UpperCAmelCase    (  ):



	'''simple docstring'''



	assert ReadInstruction("train"				,					to=1_0				,					unit="abs"				).to_spec() == "train[:10]"
	assert ReadInstruction("train"				,					from_=-8_0				,					to=1_0				,					unit="%"				).to_spec() == "train[-80%:10%]"

	UpperCAmelCase:					str 						=     "train+test"
	assert ReadInstruction.from_spec(snake_case_				).to_spec() == spec_train_test

	UpperCAmelCase:					List[Any] 						=     "train[2:10]"
	assert ReadInstruction.from_spec(snake_case_				).to_spec() == spec_train_abs

	UpperCAmelCase:					Tuple 						=     "train[15%:-20%]"
	assert ReadInstruction.from_spec(snake_case_				).to_spec() == spec_train_pct

	UpperCAmelCase:					Dict 						=     "train[:10%](closest)"
	assert ReadInstruction.from_spec(snake_case_				).to_spec() == "train[:10%]"

	UpperCAmelCase:					List[str] 						=     "train[:10%](pct1_dropremainder)"
	assert ReadInstruction.from_spec(snake_case_				).to_spec() == spec_train_pct_rounding

	UpperCAmelCase:					Any 						=     "train[:10%](pct1_dropremainder)+test[-10%:](pct1_dropremainder)"
	assert ReadInstruction.from_spec(snake_case_				).to_spec() == spec_train_test_pct_rounding



def      __UpperCAmelCase    (  ):



	'''simple docstring'''



	UpperCAmelCase:					Dict 						=     "dummy"
	UpperCAmelCase:					int 						=     [SplitInfo(name="train"				,					num_examples=1_0_0				)]
	UpperCAmelCase:					List[str] 						=     "train[:33%]"
	UpperCAmelCase:					Optional[int] 						=     "arrow"
	UpperCAmelCase:					List[Any] 						=     "prefix"

	UpperCAmelCase:					List[Any] 						=     make_file_instructions(snake_case_				,					snake_case_				,					snake_case_				,					snake_case_				,					snake_case_				)
	assert isinstance(snake_case_				,					snake_case_				)
	assert file_instructions.num_examples == 3_3
	assert file_instructions.file_instructions == [
	    {"filename": os.path.join(snake_case_				,					F'{name}-train.arrow'				), "skip": 0, "take": 3_3}
	]

	UpperCAmelCase:					int 						=     [SplitInfo(name="train"				,					num_examples=1_0_0				,					shard_lengths=[1_0] * 1_0				)]
	UpperCAmelCase:					List[Any] 						=     make_file_instructions(snake_case_				,					snake_case_				,					snake_case_				,					snake_case_				,					snake_case_				)
	assert isinstance(snake_case_				,					snake_case_				)
	assert file_instructions.num_examples == 3_3
	assert file_instructions.file_instructions == [
	    {"filename": os.path.join(snake_case_				,					F'{name}-train-00000-of-00010.arrow'				), "skip": 0, "take": -1},
	    {"filename": os.path.join(snake_case_				,					F'{name}-train-00001-of-00010.arrow'				), "skip": 0, "take": -1},
	    {"filename": os.path.join(snake_case_				,					F'{name}-train-00002-of-00010.arrow'				), "skip": 0, "take": -1},
	    {"filename": os.path.join(snake_case_				,					F'{name}-train-00003-of-00010.arrow'				), "skip": 0, "take": 3},
	]



@pytest.mark.parametrize("name, expected_exception"				,					[(None, TypeError), ("", ValueError)]				)
def      __UpperCAmelCase    (  snake_case_							:		Union[str, Any]				,					snake_case_							:		Dict				):



	'''simple docstring'''



	UpperCAmelCase:					Optional[int] 						=     [SplitInfo(name="train"				,					num_examples=1_0_0				)]
	UpperCAmelCase:					Union[str, Any] 						=     "train"
	UpperCAmelCase:					List[str] 						=     "arrow"
	UpperCAmelCase:					Any 						=     "prefix_path"
	with pytest.raises(snake_case_				):
		UpperCAmelCase:					List[Any] 						=     make_file_instructions(snake_case_				,					snake_case_				,					snake_case_				,					snake_case_				,					snake_case_				)
