


import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


__UpperCamelCase				:  Optional[int]      					=							datasets.utils.logging.get_logger(__name__)





@dataclass
class  lowerCAmelCase__(    datasets.BuilderConfig						):






    '''simple docstring'''





    A_  :     Optional[datasets.Features]			= None
    A_  :     str			= "utf-8"
    A_  :     Optional[str]			= None
    A_  :     Optional[str]			= None
    A_  :     bool			= True  # deprecated
    A_  :     Optional[int]			= None  # deprecated
    A_  :     int			= 1_0 << 2_0  # 10MB
    A_  :     Optional[bool]			= None





class  lowerCAmelCase__(    datasets.ArrowBasedBuilder						):






    '''simple docstring'''





    A_  :     Optional[int]			= JsonConfig





    def 					_lowerCamelCase (	self				:   Union[str, Any]	):





       '''simple docstring'''
       if self.config.block_size is not None:
          logger.warning('''The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead'''	)
          UpperCAmelCase_							:      Any						=   self.config.block_size
       if self.config.use_threads is not True:
          logger.warning(
              '''The JSON loader parameter `use_threads` is deprecated and doesn\'t have any effect anymore.'''	)
       if self.config.newlines_in_values is not None:
          raise ValueError('''The JSON loader parameter `newlines_in_values` is no longer supported'''	)
       return datasets.DatasetInfo(features=self.config.features	)





    def 					_lowerCamelCase (	self				:   Tuple       ,   __snake_case				:   str	):





       '''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}'''	)
       UpperCAmelCase_							:      int						=   dl_manager.download_and_extract(self.config.data_files	)
       if isinstance(__snake_case       ,   (str, list, tuple)	):
          UpperCAmelCase_							:      Optional[Any]						=   data_files
          if isinstance(__snake_case       ,   __snake_case	):
             UpperCAmelCase_							:      Any						=   [files]
          UpperCAmelCase_							:      Any						=   [dl_manager.iter_files(__snake_case	) for file in files]
          return [datasets.SplitGenerator(name=datasets.Split.TRAIN       ,   gen_kwargs={'''files''': files}	)]
       UpperCAmelCase_							:      Optional[int]						=   []
       for split_name, files in data_files.items():
          if isinstance(__snake_case       ,   __snake_case	):
             UpperCAmelCase_							:      int						=   [files]
          UpperCAmelCase_							:      Optional[Any]						=   [dl_manager.iter_files(__snake_case	) for file in files]
          splits.append(datasets.SplitGenerator(name=__snake_case       ,   gen_kwargs={'''files''': files}	)	)
       return splits





    def 					_lowerCamelCase (	self				:   Optional[int]       ,   __snake_case				:   pa.Table	):





       '''simple docstring'''
       if self.config.features is not None:
          # adding missing columns
          for column_name in set(self.config.features	) - set(pa_table.column_names	):
             UpperCAmelCase_							:      Any						=   self.config.features.arrow_schema.field(__snake_case	).type
             UpperCAmelCase_							:      Optional[int]						=   pa_table.append_column(__snake_case       ,   pa.array([None] * len(__snake_case	)       ,   type=__snake_case	)	)
          # more expensive cast to support nested structures with keys in a different order
          # allows str <-> int/float or str to Audio for example
          UpperCAmelCase_							:      Union[str, Any]						=   table_cast(__snake_case       ,   self.config.features.arrow_schema	)
       return pa_table





    def 					_lowerCamelCase (	self				:   Any       ,   __snake_case				:   Union[str, Any]	):





       '''simple docstring'''
       for file_idx, file in enumerate(itertools.chain.from_iterable(__snake_case	)	):
          # If the file is one json object and if we need to look at the list of items in one specific field
          if self.config.field is not None:
             with open(__snake_case       ,   encoding=self.config.encoding       ,   errors=self.config.encoding_errors	) as f:
                UpperCAmelCase_							:      Optional[Any]						=   json.load(__snake_case	)

             # We keep only the field we are interested in
             UpperCAmelCase_							:      List[Any]						=   dataset[self.config.field]

             # We accept two format: a list of dicts or a dict of lists
             if isinstance(__snake_case       ,   (list, tuple)	):
                UpperCAmelCase_							:      Optional[Any]						=   set().union(*[row.keys() for row in dataset]	)
                UpperCAmelCase_							:      int						=   {col: [row.get(__snake_case	) for row in dataset] for col in keys}
             else:
                UpperCAmelCase_							:      List[Any]						=   dataset
             UpperCAmelCase_							:      Any						=   pa.Table.from_pydict(__snake_case	)
             yield file_idx, self._cast_table(__snake_case	)

          # If the file has one json object per line
          else:
             with open(__snake_case       ,   '''rb'''	) as f:
                UpperCAmelCase_							:      Dict						=   0
                # Use block_size equal to the chunk size divided by 32 to leverage multithreading
                # Set a default minimum value of 16kB if the chunk size is really small
                UpperCAmelCase_							:      int						=   max(self.config.chunksize // 32       ,   16 << 10	)
                UpperCAmelCase_							:      int						=   (
                    self.config.encoding_errors if self.config.encoding_errors is not None else '''strict'''
                )
                while True:
                   UpperCAmelCase_							:      List[Any]						=   f.read(self.config.chunksize	)
                   if not batch:
                      break
                   # Finish current line
                   try:
                      batch += f.readline()
                   except (AttributeError, io.UnsupportedOperation):
                      batch += readline(__snake_case	)
                   # PyArrow only accepts utf-8 encoded bytes
                   if self.config.encoding != "utf-8":
                      UpperCAmelCase_							:      Dict						=   batch.decode(self.config.encoding       ,   errors=__snake_case	).encode('''utf-8'''	)
                   try:
                      while True:
                         try:
                            UpperCAmelCase_							:      Optional[int]						=   paj.read_json(
                                io.BytesIO(__snake_case	)       ,   read_options=paj.ReadOptions(block_size=__snake_case	)	)
                            break
                         except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
                            if (
                                isinstance(__snake_case       ,   pa.ArrowInvalid	)
                                and "straddling" not in str(__snake_case	)
                                or block_size > len(__snake_case	)
                            ):
                               raise
                            else:
                               # Increase the block size in case it was too small.
                               # The block size will be reset for the next file.
                               logger.debug(
                                   f'''Batch of {len(__snake_case	)} bytes couldn\'t be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}.'''	)
                               block_size *= 2
                   except pa.ArrowInvalid as e:
                      try:
                         with open(
                             __snake_case       ,   encoding=self.config.encoding       ,   errors=self.config.encoding_errors	) as f:
                            UpperCAmelCase_							:      Union[str, Any]						=   json.load(__snake_case	)
                      except json.JSONDecodeError:
                         logger.error(f'''Failed to read file \'{file}\' with error {type(__snake_case	)}: {e}'''	)
                         raise e
                      # If possible, parse the file as a list of json objects and exit the loop
                      if isinstance(__snake_case       ,   __snake_case	):  # list is the only sequence type supported in JSON
                         try:
                            UpperCAmelCase_							:      Union[str, Any]						=   set().union(*[row.keys() for row in dataset]	)
                            UpperCAmelCase_							:      List[Any]						=   {col: [row.get(__snake_case	) for row in dataset] for col in keys}
                            UpperCAmelCase_							:      Any						=   pa.Table.from_pydict(__snake_case	)
                         except (pa.ArrowInvalid, AttributeError) as e:
                            logger.error(f'''Failed to read file \'{file}\' with error {type(__snake_case	)}: {e}'''	)
                            raise ValueError(f'''Not able to read records in the JSON file at {file}.'''	) from None
                         yield file_idx, self._cast_table(__snake_case	)
                         break
                      else:
                         logger.error(f'''Failed to read file \'{file}\' with error {type(__snake_case	)}: {e}'''	)
                         raise ValueError(
                             f'''Not able to read records in the JSON file at {file}. '''
                             f'''You should probably indicate the field of the JSON file containing your records. '''
                             f'''This JSON file contain the following fields: {str(list(dataset.keys()	)	)}. '''
                             f'''Select the correct one and provide it as `field=\'XXX\'` to the dataset loading method. '''	) from None
                        # 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 (file_idx, batch_idx), self._cast_table(__snake_case	)
                   batch_idx += 1



# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


__UpperCamelCase				:  List[Any]      					=							{
    'configuration_vivit': ['VIVIT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'VivitConfig'],
}
try:
       if not is_vision_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __UpperCamelCase				:  Optional[Any]      					=							['VivitImageProcessor']


try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __UpperCamelCase				:  List[Any]      					=							[
           'VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST',
           'VivitModel',
           'VivitPreTrainedModel',
           'VivitForVideoClassification',
       ]


if TYPE_CHECKING:
       from .configuration_vivit import VIVIT_PRETRAINED_CONFIG_ARCHIVE_MAP, VivitConfig

       try:
              if not is_vision_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .image_processing_vivit import VivitImageProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_vivit import (
                  VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
                  VivitForVideoClassification,
                  VivitModel,
                  VivitPreTrainedModel,
              )


else:
       import sys

       __UpperCamelCase				:  str      					=							_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)