



'''simple docstring'''






import itertools
from dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa
import pyarrow.parquet as pq

import datasets
from datasets.table import table_cast


UpperCamelCase__:   int								=			datasets.utils.logging.get_logger(__name__)

@dataclass
class 				SCREAMING_SNAKE_CASE(							datasets.BuilderConfig			):







	"""simple docstring"""






	lowerCamelCase__			 =       10_000
	lowerCamelCase__			 =       None
	lowerCamelCase__			 =       None

class 				SCREAMING_SNAKE_CASE(							datasets.ArrowBasedBuilder			):







	"""simple docstring"""






	lowerCamelCase__			 =       ParquetConfig







	def      A	(						self			:     Union[str, Any]		)			->		str:
								return datasets.DatasetInfo(features=self.config.features		)







	def      A	(						self			:     str		,			__snake_case			:     Dict		)			->		Optional[int]:
								if not self.config.data_files:
															raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}"""		)
								UpperCAmelCase :       Dict 							=		dl_manager.download_and_extract(self.config.data_files		)
								if isinstance(__snake_case		,			(str, list, tuple)		):
															UpperCAmelCase :       List[str] 							=		data_files
															if isinstance(__snake_case		,			__snake_case		):
																						UpperCAmelCase :       Optional[int] 							=		[files]
															# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
															UpperCAmelCase :       int 							=		[dl_manager.iter_files(__snake_case		) for file in files]
															return [datasets.SplitGenerator(name=datasets.Split.TRAIN		,			gen_kwargs={'''files''': files}		)]
								UpperCAmelCase :       Union[str, Any] 							=		[]
								for split_name, files in data_files.items():
															if isinstance(__snake_case		,			__snake_case		):
																						UpperCAmelCase :       List[Any] 							=		[files]
															# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
															UpperCAmelCase :       List[Any] 							=		[dl_manager.iter_files(__snake_case		) for file in files]
															# Infer features is they are stoed in the arrow schema
															if self.info.features is None:
																						for file in itertools.chain.from_iterable(__snake_case		):
																													with open(__snake_case		,			'''rb'''		) as f:
																																				UpperCAmelCase :       Optional[Any] 							=		datasets.Features.from_arrow_schema(pq.read_schema(__snake_case		)		)
																													break
															splits.append(datasets.SplitGenerator(name=__snake_case		,			gen_kwargs={'''files''': files}		)		)
								return splits







	def      A	(						self			:     List[str]		,			__snake_case			:     pa.Table		)			->		pa.Table:
								if self.info.features is not None:
															# more expensive cast to support nested features with keys in a different order
															# allows str <-> int/float or str to Audio for example
															UpperCAmelCase :       Any 							=		table_cast(__snake_case		,			self.info.features.arrow_schema		)
								return pa_table







	def      A	(						self			:     List[Any]		,			__snake_case			:     Tuple		)			->		Dict:
								UpperCAmelCase :       List[Any] 							=		self.info.features.arrow_schema if self.info.features is not None else None
								if self.info.features is not None and self.config.columns is not None:
															if sorted(field.name for field in schema		) != sorted(self.config.columns		):
																						raise ValueError(
																						    F"""Tried to load parquet data with columns '{self.config.columns}' with mismatching features '{self.info.features}'"""		)
								for file_idx, file in enumerate(itertools.chain.from_iterable(__snake_case		)		):
															with open(__snake_case		,			'''rb'''		) as f:
																						UpperCAmelCase :       Union[str, Any] 							=		pq.ParquetFile(__snake_case		)
																						try:
																													for batch_idx, record_batch in enumerate(
																													    parquet_file.iter_batches(batch_size=self.config.batch_size		,			columns=self.config.columns		)		):
																																				UpperCAmelCase :       List[Any] 							=		pa.Table.from_batches([record_batch]		)
																																				# Uncomment for debugging (will print the Arrow table size and elements)
																																				# logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
																																				# logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																																				yield F"""{file_idx}_{batch_idx}""", self._cast_table(__snake_case		)
																						except ValueError as e:
																													logger.error(F"""Failed to read file '{file}' with error {type(__snake_case		)}: {e}"""		)
																													raise










'''simple docstring'''






import unittest

import numpy as np

from transformers.testing_utils import require_flax, require_tf, require_torch
from transformers.utils import (
    expand_dims,
    flatten_dict,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    reshape,
    squeeze,
    transpose,
)


if is_flax_available():
						import jax.numpy as jnp

if is_tf_available():
						import tensorflow as tf

if is_torch_available():
						import torch

class 				SCREAMING_SNAKE_CASE(							unittest.TestCase			):







	"""simple docstring"""






	def      A	(						self			:     Any		)			->		str:
								UpperCAmelCase :       Any 							=		{
								    '''task_specific_params''': {
								        '''summarization''': {'''length_penalty''': 1.0, '''max_length''': 128, '''min_length''': 12, '''num_beams''': 4},
								        '''summarization_cnn''': {'''length_penalty''': 2.0, '''max_length''': 142, '''min_length''': 56, '''num_beams''': 4},
								        '''summarization_xsum''': {'''length_penalty''': 1.0, '''max_length''': 62, '''min_length''': 11, '''num_beams''': 6},
								    }
								}
								UpperCAmelCase :       int 							=		{
								    '''task_specific_params.summarization.length_penalty''': 1.0,
								    '''task_specific_params.summarization.max_length''': 128,
								    '''task_specific_params.summarization.min_length''': 12,
								    '''task_specific_params.summarization.num_beams''': 4,
								    '''task_specific_params.summarization_cnn.length_penalty''': 2.0,
								    '''task_specific_params.summarization_cnn.max_length''': 142,
								    '''task_specific_params.summarization_cnn.min_length''': 56,
								    '''task_specific_params.summarization_cnn.num_beams''': 4,
								    '''task_specific_params.summarization_xsum.length_penalty''': 1.0,
								    '''task_specific_params.summarization_xsum.max_length''': 62,
								    '''task_specific_params.summarization_xsum.min_length''': 11,
								    '''task_specific_params.summarization_xsum.num_beams''': 6,
								}

								self.assertEqual(flatten_dict(__snake_case		)		,			__snake_case		)







	def      A	(						self			:     int		)			->		str:
								UpperCAmelCase :       Any 							=		np.random.randn(3		,			4		)
								self.assertTrue(np.allclose(transpose(__snake_case		)		,			x.transpose()		)		)

								UpperCAmelCase :       str 							=		np.random.randn(3		,			4		,			5		)
								self.assertTrue(np.allclose(transpose(__snake_case		,			axes=(1, 2, 0)		)		,			x.transpose((1, 2, 0)		)		)		)







	@require_torch
	def      A	(						self			:     str		)			->		Union[str, Any]:
								UpperCAmelCase :       Any 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       List[Any] 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(transpose(__snake_case		)		,			transpose(__snake_case		).numpy()		)		)

								UpperCAmelCase :       Tuple 							=		np.random.randn(3		,			4		,			5		)
								UpperCAmelCase :       Any 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(transpose(__snake_case		,			axes=(1, 2, 0)		)		,			transpose(__snake_case		,			axes=(1, 2, 0)		).numpy()		)		)







	@require_tf
	def      A	(						self			:     List[str]		)			->		Optional[Any]:
								UpperCAmelCase :       int 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       Optional[int] 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(transpose(__snake_case		)		,			transpose(__snake_case		).numpy()		)		)

								UpperCAmelCase :       Union[str, Any] 							=		np.random.randn(3		,			4		,			5		)
								UpperCAmelCase :       str 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(transpose(__snake_case		,			axes=(1, 2, 0)		)		,			transpose(__snake_case		,			axes=(1, 2, 0)		).numpy()		)		)







	@require_flax
	def      A	(						self			:     Tuple		)			->		Any:
								UpperCAmelCase :       List[Any] 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       List[str] 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(transpose(__snake_case		)		,			np.asarray(transpose(__snake_case		)		)		)		)

								UpperCAmelCase :       Dict 							=		np.random.randn(3		,			4		,			5		)
								UpperCAmelCase :       int 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(transpose(__snake_case		,			axes=(1, 2, 0)		)		,			np.asarray(transpose(__snake_case		,			axes=(1, 2, 0)		)		)		)		)







	def      A	(						self			:     Optional[Any]		)			->		Any:
								UpperCAmelCase :       Union[str, Any] 							=		np.random.randn(3		,			4		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(4, 3)		)		,			np.reshape(__snake_case		,			(4, 3)		)		)		)

								UpperCAmelCase :       Union[str, Any] 							=		np.random.randn(3		,			4		,			5		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(12, 5)		)		,			np.reshape(__snake_case		,			(12, 5)		)		)		)







	@require_torch
	def      A	(						self			:     Union[str, Any]		)			->		int:
								UpperCAmelCase :       Dict 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       Optional[Any] 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(4, 3)		)		,			reshape(__snake_case		,			(4, 3)		).numpy()		)		)

								UpperCAmelCase :       List[str] 							=		np.random.randn(3		,			4		,			5		)
								UpperCAmelCase :       List[Any] 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(12, 5)		)		,			reshape(__snake_case		,			(12, 5)		).numpy()		)		)







	@require_tf
	def      A	(						self			:     int		)			->		List[str]:
								UpperCAmelCase :       List[Any] 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       List[str] 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(4, 3)		)		,			reshape(__snake_case		,			(4, 3)		).numpy()		)		)

								UpperCAmelCase :       List[Any] 							=		np.random.randn(3		,			4		,			5		)
								UpperCAmelCase :       Optional[Any] 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(12, 5)		)		,			reshape(__snake_case		,			(12, 5)		).numpy()		)		)







	@require_flax
	def      A	(						self			:     Any		)			->		Dict:
								UpperCAmelCase :       Tuple 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       Union[str, Any] 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(4, 3)		)		,			np.asarray(reshape(__snake_case		,			(4, 3)		)		)		)		)

								UpperCAmelCase :       Any 							=		np.random.randn(3		,			4		,			5		)
								UpperCAmelCase :       Optional[Any] 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(12, 5)		)		,			np.asarray(reshape(__snake_case		,			(12, 5)		)		)		)		)







	def      A	(						self			:     List[Any]		)			->		List[Any]:
								UpperCAmelCase :       Union[str, Any] 							=		np.random.randn(1		,			3		,			4		)
								self.assertTrue(np.allclose(squeeze(__snake_case		)		,			np.squeeze(__snake_case		)		)		)

								UpperCAmelCase :       str 							=		np.random.randn(1		,			4		,			1		,			5		)
								self.assertTrue(np.allclose(squeeze(__snake_case		,			axis=2		)		,			np.squeeze(__snake_case		,			axis=2		)		)		)







	@require_torch
	def      A	(						self			:     Optional[int]		)			->		Union[str, Any]:
								UpperCAmelCase :       Optional[int] 							=		np.random.randn(1		,			3		,			4		)
								UpperCAmelCase :       List[str] 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(squeeze(__snake_case		)		,			squeeze(__snake_case		).numpy()		)		)

								UpperCAmelCase :       Any 							=		np.random.randn(1		,			4		,			1		,			5		)
								UpperCAmelCase :       str 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(squeeze(__snake_case		,			axis=2		)		,			squeeze(__snake_case		,			axis=2		).numpy()		)		)







	@require_tf
	def      A	(						self			:     Optional[Any]		)			->		Dict:
								UpperCAmelCase :       int 							=		np.random.randn(1		,			3		,			4		)
								UpperCAmelCase :       Optional[int] 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(squeeze(__snake_case		)		,			squeeze(__snake_case		).numpy()		)		)

								UpperCAmelCase :       List[str] 							=		np.random.randn(1		,			4		,			1		,			5		)
								UpperCAmelCase :       Optional[int] 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(squeeze(__snake_case		,			axis=2		)		,			squeeze(__snake_case		,			axis=2		).numpy()		)		)







	@require_flax
	def      A	(						self			:     List[Any]		)			->		Dict:
								UpperCAmelCase :       Optional[int] 							=		np.random.randn(1		,			3		,			4		)
								UpperCAmelCase :       int 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(squeeze(__snake_case		)		,			np.asarray(squeeze(__snake_case		)		)		)		)

								UpperCAmelCase :       str 							=		np.random.randn(1		,			4		,			1		,			5		)
								UpperCAmelCase :       int 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(squeeze(__snake_case		,			axis=2		)		,			np.asarray(squeeze(__snake_case		,			axis=2		)		)		)		)







	def      A	(						self			:     Optional[Any]		)			->		int:
								UpperCAmelCase :       Optional[Any] 							=		np.random.randn(3		,			4		)
								self.assertTrue(np.allclose(expand_dims(__snake_case		,			axis=1		)		,			np.expand_dims(__snake_case		,			axis=1		)		)		)







	@require_torch
	def      A	(						self			:     List[str]		)			->		Tuple:
								UpperCAmelCase :       Tuple 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       Tuple 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(expand_dims(__snake_case		,			axis=1		)		,			expand_dims(__snake_case		,			axis=1		).numpy()		)		)







	@require_tf
	def      A	(						self			:     List[str]		)			->		Union[str, Any]:
								UpperCAmelCase :       Union[str, Any] 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       Any 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(expand_dims(__snake_case		,			axis=1		)		,			expand_dims(__snake_case		,			axis=1		).numpy()		)		)







	@require_flax
	def      A	(						self			:     Any		)			->		List[Any]:
								UpperCAmelCase :       List[str] 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       str 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(expand_dims(__snake_case		,			axis=1		)		,			np.asarray(expand_dims(__snake_case		,			axis=1		)		)		)		)





