
"""simple docstring"""
from argparse import ArgumentParser

from . import BaseTransformersCLICommand





def    _lowerCAmelCase    (     lowerCamelCase__							:							Dict				)  ->    Optional[int]:
  return DownloadCommand(args.model,						args.cache_dir,						args.force,						args.trust_remote_code				)





class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE  ):
     @staticmethod
     def 				__SCREAMING_SNAKE_CASE							(     snake_case__				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Dict							    = parser.add_parser("download"				)
       download_parser.add_argument(
           "--cache-dir"  ,     type=snake_case__  ,     default=snake_case__  ,     help="Path to location to store the models"				)
       download_parser.add_argument(
           "--force"  ,     action="store_true"  ,     help="Force the model to be download even if already in cache-dir"				)
       download_parser.add_argument(
           "--trust-remote-code"  ,     action="store_true"  ,     help="Whether or not to allow for custom models defined on the Hub in their own modeling files. Use only if you've reviewed the code as it will execute on your local machine"  ,     )
       download_parser.add_argument("model"  ,     type=snake_case__  ,     help="Name of the model to download"				)
       download_parser.set_defaults(func=snake_case__				)
     def __init__(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				int							    = model
       _SCREAMING_SNAKE_CASE   :				int							    = cache
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = force
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = trust_remote_code




     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       from ..models.auto import AutoModel, AutoTokenizer

       AutoModel.from_pretrained(
           self._model  ,     cache_dir=self._cache  ,     force_download=self._force  ,     trust_remote_code=self._trust_remote_code				)
       AutoTokenizer.from_pretrained(
           self._model  ,     cache_dir=self._cache  ,     force_download=self._force  ,     trust_remote_code=self._trust_remote_code				)





"""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


lowercase_       :						List[str]     					=	datasets.utils.logging.get_logger(__name__)
@dataclass
class 			UpperCamelCase      (							datasets.BuilderConfig  ):
     A__      =			10000
     A__      =			None
     A__      =			None
class 			UpperCamelCase      (							datasets.ArrowBasedBuilder  ):
     A__      =			ParquetConfig
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       return datasets.DatasetInfo(features=self.config.features				)
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__				):




       """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}'''				)
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = dl_manager.download_and_extract(self.config.data_files				)
       if isinstance(snake_case__  ,     (str, list, tuple)				):
         _SCREAMING_SNAKE_CASE   :				int							    = data_files
         if isinstance(snake_case__  ,     snake_case__				):
           _SCREAMING_SNAKE_CASE   :				List[Any]							    = [files]
         # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
         _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = [dl_manager.iter_files(snake_case__				) for file in files]
         return [datasets.SplitGenerator(name=datasets.Split.TRAIN  ,     gen_kwargs={"files": files}				)]
       _SCREAMING_SNAKE_CASE   :				Tuple							    = []
       for split_name, files in data_files.items():
         if isinstance(snake_case__  ,     snake_case__				):
           _SCREAMING_SNAKE_CASE   :				Optional[int]							    = [files]
         # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
         _SCREAMING_SNAKE_CASE   :				Tuple							    = [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:
               _SCREAMING_SNAKE_CASE   :				Optional[int]							    = 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 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__				):




       """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
         _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = table_cast(snake_case__  ,     self.info.features.arrow_schema				)
       return pa_table




     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				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:
           _SCREAMING_SNAKE_CASE   :				List[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				)				):
               _SCREAMING_SNAKE_CASE   :				int							    = 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



