


import itertools
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa

import datasets
from datasets.table import table_cast


_snake_case					=					datasets.utils.logging.get_logger(__name__)
@dataclass
class      UpperCAmelCase_ ( datasets.BuilderConfig   ):

			'''simple docstring'''




			__A			:				Optional[datasets.Features]							=		None



class      UpperCAmelCase_ ( datasets.ArrowBasedBuilder   ):

			'''simple docstring'''




			__A			:				str							=		ArrowConfig






			def _snake_case   (		self			):






					"""simple docstring"""





					return datasets.DatasetInfo(features=self.config.features			)






			def _snake_case   (		self						,    __A			):






					"""simple docstring"""





					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}"""			)
					lowerCamelCase			: List[Any]   =    dl_manager.download_and_extract(self.config.data_files			)
					if isinstance(__A						,    (str, list, tuple)			):
							lowerCamelCase			: List[Any]   =    data_files
							if isinstance(__A						,    __A			):
									lowerCamelCase			: Any   =    [files]
							# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
							lowerCamelCase			: List[str]   =    [dl_manager.iter_files(__A			) for file in files]
							return [datasets.SplitGenerator(name=datasets.Split.TRAIN						,    gen_kwargs={"files": files}			)]
					lowerCamelCase			: Union[str, Any]   =    []
					for split_name, files in data_files.items():
							if isinstance(__A						,    __A			):
									lowerCamelCase			: Optional[Any]   =    [files]
							# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
							lowerCamelCase			: Dict   =    [dl_manager.iter_files(__A			) 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(__A			):
											with open(__A						,    "rb"			) as f:
													lowerCamelCase			: List[Any]   =    datasets.Features.from_arrow_schema(pa.ipc.open_stream(__A			).schema			)
											break
							splits.append(datasets.SplitGenerator(name=__A						,    gen_kwargs={"files": files}			)			)
					return splits






			def _snake_case   (		self						,    __A			):






					"""simple docstring"""





					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
							lowerCamelCase			: int   =    table_cast(__A						,    self.info.features.arrow_schema			)
					return pa_table






			def _snake_case   (		self						,    __A			):






					"""simple docstring"""





					for file_idx, file in enumerate(itertools.chain.from_iterable(__A			)			):
							with open(__A						,    "rb"			) as f:
									try:
											for batch_idx, record_batch in enumerate(pa.ipc.open_stream(__A			)			):
													lowerCamelCase			: Union[str, 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(__A			)
									except ValueError as e:
											logger.error(F"""Failed to read file '{file}' with error {type(__A			)}: {e}"""			)
											raise







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


_snake_case					=					{
    '''configuration_rwkv''': ['''RWKV_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''RwkvConfig''', '''RwkvOnnxConfig'''],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _snake_case					=					[
       '''RWKV_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''RwkvForCausalLM''',
       '''RwkvModel''',
       '''RwkvPreTrainedModel''',
   ]


if TYPE_CHECKING:
   from .configuration_rwkv import RWKV_PRETRAINED_CONFIG_ARCHIVE_MAP, RwkvConfig, RwkvOnnxConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_rwkv import (
          RWKV_PRETRAINED_MODEL_ARCHIVE_LIST,
          RwkvForCausalLM,
          RwkvModel,
          RwkvPreTrainedModel,
      )
else:
   import sys

   _snake_case					=					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



