





"""simple docstring"""




import io
import json

import fsspec
import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.json import JsonDatasetReader, JsonDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases





def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase					):






  """simple docstring"""





  assert isinstance(_UpperCAmelCase ,					_UpperCAmelCase					)
  assert dataset.num_rows == 4
  assert dataset.num_columns == 3
  assert dataset.column_names == ["col_1", "col_2", "col_3"]
  for feature, expected_dtype in expected_features.items():
    assert dataset.features[feature].dtype == expected_dtype





@pytest.mark.parametrize('keep_in_memory' ,					[False, True]					)
def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase					):






  """simple docstring"""





  A_  :       Union[str, Any]										=  tmp_path / 'cache'
  A_  :       Tuple										=  {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
  with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
    A_  :       int										=  JsonDatasetReader(_UpperCAmelCase ,					cache_dir=_UpperCAmelCase ,					keep_in_memory=_UpperCAmelCase					).read()
  _check_json_dataset(_UpperCAmelCase ,					_UpperCAmelCase					)





@pytest.mark.parametrize(
    'features' ,					[
        None,
        {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'},
        {'col_1': 'string', 'col_2': 'string', 'col_3': 'string'},
        {'col_1': 'int32', 'col_2': 'int32', 'col_3': 'int32'},
        {'col_1': 'float32', 'col_2': 'float32', 'col_3': 'float32'},
    ] ,					)
def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase					):






  """simple docstring"""





  A_  :       Union[str, Any]										=  tmp_path / 'cache'
  A_  :       Union[str, Any]										=  {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
  A_  :       int										=  features.copy() if features else default_expected_features
  A_  :       List[Any]										=  (
      Features({feature: Value(_UpperCAmelCase					) for feature, dtype in features.items()}					) if features is not None else None
  )
  A_  :       str										=  JsonDatasetReader(_UpperCAmelCase ,					features=_UpperCAmelCase ,					cache_dir=_UpperCAmelCase					).read()
  _check_json_dataset(_UpperCAmelCase ,					_UpperCAmelCase					)





@pytest.mark.parametrize(
    'features' ,					[
        None,
        {'col_3': 'float64', 'col_1': 'string', 'col_2': 'int64'},
    ] ,					)
def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase					):






  """simple docstring"""





  A_  :       Dict										=  tmp_path / 'cache'
  A_  :       Tuple										=  {'col_3': 'float64', 'col_1': 'string', 'col_2': 'int64'}
  A_  :       List[Any]										=  features.copy() if features else default_expected_features
  A_  :       Dict										=  (
      Features({feature: Value(_UpperCAmelCase					) for feature, dtype in features.items()}					) if features is not None else None
  )
  A_  :       List[str]										=  JsonDatasetReader(_UpperCAmelCase ,					features=_UpperCAmelCase ,					cache_dir=_UpperCAmelCase					).read()
  assert isinstance(_UpperCAmelCase ,					_UpperCAmelCase					)
  assert dataset.num_rows == 2
  assert dataset.num_columns == 3
  assert dataset.column_names == ["col_3", "col_1", "col_2"]
  for feature, expected_dtype in expected_features.items():
    assert dataset.features[feature].dtype == expected_dtype





def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase					):






  """simple docstring"""





  A_  :       int										=  {'col_2': 'int64', 'col_3': 'float64', 'col_1': 'string'}
  A_  :       Tuple										=  features.copy()
  A_  :       Tuple										=  (
      Features({feature: Value(_UpperCAmelCase					) for feature, dtype in features.items()}					) if features is not None else None
  )
  A_  :       Optional[int]										=  tmp_path / 'cache'
  A_  :       int										=  JsonDatasetReader(_UpperCAmelCase ,					features=_UpperCAmelCase ,					cache_dir=_UpperCAmelCase					).read()
  assert isinstance(_UpperCAmelCase ,					_UpperCAmelCase					)
  assert dataset.num_rows == 2
  assert dataset.num_columns == 3
  assert dataset.column_names == ["col_2", "col_3", "col_1"]
  for feature, expected_dtype in expected_features.items():
    assert dataset.features[feature].dtype == expected_dtype





@pytest.mark.parametrize('split' ,					[None, NamedSplit('train'					), 'train', 'test']					)
def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase					):






  """simple docstring"""





  A_  :       Union[str, Any]										=  tmp_path / 'cache'
  A_  :       int										=  {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
  A_  :       Optional[Any]										=  JsonDatasetReader(_UpperCAmelCase ,					cache_dir=_UpperCAmelCase ,					split=_UpperCAmelCase					).read()
  _check_json_dataset(_UpperCAmelCase ,					_UpperCAmelCase					)
  assert dataset.split == split if split else "train"





@pytest.mark.parametrize('path_type' ,					[str, list]					)
def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase					):






  """simple docstring"""





  if issubclass(_UpperCAmelCase ,					_UpperCAmelCase					):
    A_  :       int										=  jsonl_path
  elif issubclass(_UpperCAmelCase ,					_UpperCAmelCase					):
    A_  :       Any										=  [jsonl_path]
  A_  :       List[Any]										=  tmp_path / 'cache'
  A_  :       int										=  {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
  A_  :       Union[str, Any]										=  JsonDatasetReader(_UpperCAmelCase ,					cache_dir=_UpperCAmelCase					).read()
  _check_json_dataset(_UpperCAmelCase ,					_UpperCAmelCase					)





def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase=("train",)					):






  """simple docstring"""





  assert isinstance(_UpperCAmelCase ,					_UpperCAmelCase					)
  for split in splits:
    A_  :       Optional[int]										=  dataset_dict[split]
    assert dataset.num_rows == 4
    assert dataset.num_columns == 3
    assert dataset.column_names == ["col_1", "col_2", "col_3"]
    for feature, expected_dtype in expected_features.items():
      assert dataset.features[feature].dtype == expected_dtype





@pytest.mark.parametrize('keep_in_memory' ,					[False, True]					)
def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase					):






  """simple docstring"""





  A_  :       Tuple										=  tmp_path / 'cache'
  A_  :       str										=  {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
  with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
    A_  :       Optional[Any]										=  JsonDatasetReader({'train': jsonl_path} ,					cache_dir=_UpperCAmelCase ,					keep_in_memory=_UpperCAmelCase					).read()
  _check_json_datasetdict(_UpperCAmelCase ,					_UpperCAmelCase					)





@pytest.mark.parametrize(
    'features' ,					[
        None,
        {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'},
        {'col_1': 'string', 'col_2': 'string', 'col_3': 'string'},
        {'col_1': 'int32', 'col_2': 'int32', 'col_3': 'int32'},
        {'col_1': 'float32', 'col_2': 'float32', 'col_3': 'float32'},
    ] ,					)
def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase					):






  """simple docstring"""





  A_  :       List[Any]										=  tmp_path / 'cache'
  A_  :       Optional[int]										=  {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
  A_  :       List[str]										=  features.copy() if features else default_expected_features
  A_  :       List[str]										=  (
      Features({feature: Value(_UpperCAmelCase					) for feature, dtype in features.items()}					) if features is not None else None
  )
  A_  :       Tuple										=  JsonDatasetReader({'train': jsonl_path} ,					features=_UpperCAmelCase ,					cache_dir=_UpperCAmelCase					).read()
  _check_json_datasetdict(_UpperCAmelCase ,					_UpperCAmelCase					)





@pytest.mark.parametrize('split' ,					[None, NamedSplit('train'					), 'train', 'test']					)
def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase					):






  """simple docstring"""





  if split:
    A_  :       Dict										=  {split: jsonl_path}
  else:
    A_  :       List[str]										=  'train'
    A_  :       int										=  {'train': jsonl_path, 'test': jsonl_path}
  A_  :       Any										=  tmp_path / 'cache'
  A_  :       Tuple										=  {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
  A_  :       List[Any]										=  JsonDatasetReader(_UpperCAmelCase ,					cache_dir=_UpperCAmelCase					).read()
  _check_json_datasetdict(_UpperCAmelCase ,					_UpperCAmelCase ,					splits=list(path.keys()					)					)
  assert all(dataset[split].split == split for split in path.keys()					)





def 						UpperCAmelCase__		(     _UpperCAmelCase					):






  """simple docstring"""





  return json.load(_UpperCAmelCase					)





def 						UpperCAmelCase__		(     _UpperCAmelCase					):






  """simple docstring"""





  return [json.loads(_UpperCAmelCase					) for line in buffer]





class  _UpperCAmelCase  :



      '''simple docstring'''

      @pytest.mark.parametrize('lines, load_json_function'					,      [(True, load_json_lines), (False, load_json)]  )
      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        with io.BytesIO() as buffer:
          JsonDatasetWriter(snake_case_					,      snake_case_					,      lines=snake_case_  ).write()
          buffer.seek(0  )
          A_  :       Union[str, Any]										=  load_json_function(snake_case_  )
        assert isinstance(snake_case_					,      snake_case_  )
        assert isinstance(exported_content[0]					,      snake_case_  )
        assert len(snake_case_  ) == 1_0

      @pytest.mark.parametrize(
          'orient, container, keys, len_at'					,      [
              ('records', list, {'tokens', 'labels', 'answers', 'id'}, None),
              ('split', dict, {'columns', 'data'}, 'data'),
              ('index', dict, set('0123456789'  ), None),
              ('columns', dict, {'tokens', 'labels', 'answers', 'id'}, 'tokens'),
              ('values', list, None, None),
              ('table', dict, {'schema', 'data'}, 'data'),
          ]					,      )
      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        with io.BytesIO() as buffer:
          JsonDatasetWriter(snake_case_					,      snake_case_					,      lines=snake_case_					,      orient=snake_case_  ).write()
          buffer.seek(0  )
          A_  :       str										=  load_json(snake_case_  )
        assert isinstance(snake_case_					,      snake_case_  )
        if keys:
          if container is dict:
            assert exported_content.keys() == keys
          else:
            assert exported_content[0].keys() == keys
        else:
          assert not hasattr(snake_case_					,      'keys'  ) and not hasattr(exported_content[0]					,      'keys'  )
        if len_at:
          assert len(exported_content[len_at]  ) == 1_0
        else:
          assert len(snake_case_  ) == 1_0

      @pytest.mark.parametrize('lines, load_json_function'					,      [(True, load_json_lines), (False, load_json)]  )
      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        with io.BytesIO() as buffer:
          JsonDatasetWriter(snake_case_					,      snake_case_					,      lines=snake_case_					,      num_proc=2  ).write()
          buffer.seek(0  )
          A_  :       List[Any]										=  load_json_function(snake_case_  )
        assert isinstance(snake_case_					,      snake_case_  )
        assert isinstance(exported_content[0]					,      snake_case_  )
        assert len(snake_case_  ) == 1_0

      @pytest.mark.parametrize(
          'orient, container, keys, len_at'					,      [
              ('records', list, {'tokens', 'labels', 'answers', 'id'}, None),
              ('split', dict, {'columns', 'data'}, 'data'),
              ('index', dict, set('0123456789'  ), None),
              ('columns', dict, {'tokens', 'labels', 'answers', 'id'}, 'tokens'),
              ('values', list, None, None),
              ('table', dict, {'schema', 'data'}, 'data'),
          ]					,      )
      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        with io.BytesIO() as buffer:
          JsonDatasetWriter(snake_case_					,      snake_case_					,      lines=snake_case_					,      orient=snake_case_					,      num_proc=2  ).write()
          buffer.seek(0  )
          A_  :       Optional[Any]										=  load_json(snake_case_  )
        assert isinstance(snake_case_					,      snake_case_  )
        if keys:
          if container is dict:
            assert exported_content.keys() == keys
          else:
            assert exported_content[0].keys() == keys
        else:
          assert not hasattr(snake_case_					,      'keys'  ) and not hasattr(exported_content[0]					,      'keys'  )
        if len_at:
          assert len(exported_content[len_at]  ) == 1_0
        else:
          assert len(snake_case_  ) == 1_0

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        with pytest.raises(snake_case_  ):
          with io.BytesIO() as buffer:
            JsonDatasetWriter(snake_case_					,      snake_case_					,      num_proc=0  )

      @pytest.mark.parametrize('compression, extension'					,      [('gzip', 'gz'), ('bz2', 'bz2'), ('xz', 'xz')]  )
      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        A_  :       Any										=  tmp_path_factory.mktemp('data'  ) / F"""test.json.{extension}"""
        A_  :       Any										=  str(shared_datadir / F"""test_file.json.{extension}"""  )
        JsonDatasetWriter(snake_case_					,      snake_case_					,      compression=snake_case_  ).write()

        with fsspec.open(snake_case_					,      'rb'					,      compression='infer'  ) as f:
          A_  :       Optional[int]										=  f.read()
        with fsspec.open(snake_case_					,      'rb'					,      compression='infer'  ) as f:
          A_  :       str										=  f.read()
        assert exported_content == original_content






"""simple docstring"""




import warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType

class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''
      lowercase_      : Dict	    =   ["""image_processor""", """tokenizer"""]
      lowercase_      : Union[str, Any]	    =   """ViltImageProcessor"""
      lowercase_      : Any	    =   ("""BertTokenizer""", """BertTokenizerFast""")

      def __init__(     self					,      snake_case_=None					,      snake_case_=None					,      **snake_case_  ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  None
        if "feature_extractor" in kwargs:
          warnings.warn(
              'The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'
              ' instead.'					,      snake_case_					,      )
          A_  :       Dict										=  kwargs.pop('feature_extractor'  )

        A_  :       Dict										=  image_processor if image_processor is not None else feature_extractor
        if image_processor is None:
          raise ValueError('You need to specify an `image_processor`.'  )
        if tokenizer is None:
          raise ValueError('You need to specify a `tokenizer`.'  )

        super().__init__(snake_case_					,      snake_case_  )
        A_  :       List[str]										=  self.image_processor

      def __call__(     self					,      snake_case_					,      snake_case_ = None					,      snake_case_ = True					,      snake_case_ = False					,      snake_case_ = None					,      snake_case_ = None					,      snake_case_ = 0					,      snake_case_ = None					,      snake_case_ = None					,      snake_case_ = None					,      snake_case_ = False					,      snake_case_ = False					,      snake_case_ = False					,      snake_case_ = False					,      snake_case_ = True					,      snake_case_ = None					,      **snake_case_					,      ):

        """simple docstring"""

        A_  :       str										=  self.tokenizer(
            text=snake_case_					,      add_special_tokens=snake_case_					,      padding=snake_case_					,      truncation=snake_case_					,      max_length=snake_case_					,      stride=snake_case_					,      pad_to_multiple_of=snake_case_					,      return_token_type_ids=snake_case_					,      return_attention_mask=snake_case_					,      return_overflowing_tokens=snake_case_					,      return_special_tokens_mask=snake_case_					,      return_offsets_mapping=snake_case_					,      return_length=snake_case_					,      verbose=snake_case_					,      return_tensors=snake_case_					,      **snake_case_					,      )
        # add pixel_values + pixel_mask
        A_  :       Optional[int]										=  self.image_processor(snake_case_					,      return_tensors=snake_case_  )
        encoding.update(snake_case_  )

        return encoding

      def 					lowerCamelCase_			(     self					,      *snake_case_					,      **snake_case_  ):

        """simple docstring"""

        return self.tokenizer.batch_decode(*snake_case_					,      **snake_case_  )

      def 					lowerCamelCase_			(     self					,      *snake_case_					,      **snake_case_  ):

        """simple docstring"""

        return self.tokenizer.decode(*snake_case_					,      **snake_case_  )

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Any										=  self.tokenizer.model_input_names
        A_  :       Any										=  self.image_processor.model_input_names
        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names  )  )

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        warnings.warn(
            '`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'					,      snake_case_					,      )
        return self.image_processor_class

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        warnings.warn(
            '`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'					,      snake_case_					,      )
        return self.image_processor