






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCamelCase      :       Tuple    =						{
    """configuration_bigbird_pegasus""": [
        """BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """BigBirdPegasusConfig""",
        """BigBirdPegasusOnnxConfig""",
    ],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase      :       Any    =						[
     """BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """BigBirdPegasusForCausalLM""",
     """BigBirdPegasusForConditionalGeneration""",
     """BigBirdPegasusForQuestionAnswering""",
     """BigBirdPegasusForSequenceClassification""",
     """BigBirdPegasusModel""",
     """BigBirdPegasusPreTrainedModel""",
 ]


if TYPE_CHECKING:
 from .configuration_bigbird_pegasus import (
     BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP,
     BigBirdPegasusConfig,
     BigBirdPegasusOnnxConfig,
 )

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_bigbird_pegasus import (
      BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST,
      BigBirdPegasusForCausalLM,
      BigBirdPegasusForConditionalGeneration,
      BigBirdPegasusForQuestionAnswering,
      BigBirdPegasusForSequenceClassification,
      BigBirdPegasusModel,
      BigBirdPegasusPreTrainedModel,
  )


else:
 import sys

 lowerCamelCase      :       Any    =						_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)













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 		lowerCAmelCase  (	UpperCamelCase__   :   Any       ,			UpperCamelCase__   :   Optional[Any]					)      ->    str:

						"""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 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[int]       ,			UpperCamelCase__   :   Optional[int]       ,			UpperCamelCase__   :   List[Any]					)      ->    Tuple:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Dict           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Dict           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
												__SCREAMING_SNAKE_CASE:				Tuple           =   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 		lowerCAmelCase  (	UpperCamelCase__   :   List[str]       ,			UpperCamelCase__   :   int       ,			UpperCamelCase__   :   Optional[int]					)      ->    int:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				List[str]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Tuple           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE:				Any           =   features.copy() if features else default_expected_features
						__SCREAMING_SNAKE_CASE:				List[Any]           =   (
						    Features({feature: Value(UpperCamelCase__					) for feature, dtype in features.items()}					) if features is not None else None
						)
						__SCREAMING_SNAKE_CASE:				Tuple           =   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 		lowerCAmelCase  (	UpperCamelCase__   :   Union[str, Any]       ,			UpperCamelCase__   :   List[str]       ,			UpperCamelCase__   :   Optional[int]					)      ->    Dict:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Optional[int]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Dict           =   {'''col_3''': '''float64''', '''col_1''': '''string''', '''col_2''': '''int64'''}
						__SCREAMING_SNAKE_CASE:				int           =   features.copy() if features else default_expected_features
						__SCREAMING_SNAKE_CASE:				List[Any]           =   (
						    Features({feature: Value(UpperCamelCase__					) for feature, dtype in features.items()}					) if features is not None else None
						)
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   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 		lowerCAmelCase  (	UpperCamelCase__   :   str       ,			UpperCamelCase__   :   Optional[int]					)      ->    str:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				str           =   {'''col_2''': '''int64''', '''col_3''': '''float64''', '''col_1''': '''string'''}
						__SCREAMING_SNAKE_CASE:				Dict           =   features.copy()
						__SCREAMING_SNAKE_CASE:				Optional[int]           =   (
						    Features({feature: Value(UpperCamelCase__					) for feature, dtype in features.items()}					) if features is not None else None
						)
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Tuple           =   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 		lowerCAmelCase  (	UpperCamelCase__   :   Any       ,			UpperCamelCase__   :   str       ,			UpperCamelCase__   :   int					)      ->    Optional[Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				List[Any]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE:				Dict           =   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 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[Any]       ,			UpperCamelCase__   :   Union[str, Any]       ,			UpperCamelCase__   :   str					)      ->    Any:

						"""simple docstring"""





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






def 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[int]       ,			UpperCamelCase__   :   int       ,			UpperCamelCase__   :   Union[str, Any]=("train",)					)      ->    List[str]:

						"""simple docstring"""





						assert isinstance(UpperCamelCase__       ,			UpperCamelCase__					)
						for split in splits:
												__SCREAMING_SNAKE_CASE:				Tuple           =   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 		lowerCAmelCase  (	UpperCamelCase__   :   List[Any]       ,			UpperCamelCase__   :   Union[str, Any]       ,			UpperCamelCase__   :   Tuple					)      ->    Dict:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				List[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():
												__SCREAMING_SNAKE_CASE:				Dict           =   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 		lowerCAmelCase  (	UpperCamelCase__   :   Dict       ,			UpperCamelCase__   :   Tuple       ,			UpperCamelCase__   :   List[str]					)      ->    Union[str, Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE:				str           =   features.copy() if features else default_expected_features
						__SCREAMING_SNAKE_CASE:				Optional[int]           =   (
						    Features({feature: Value(UpperCamelCase__					) for feature, dtype in features.items()}					) if features is not None else None
						)
						__SCREAMING_SNAKE_CASE:				Dict           =   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 		lowerCAmelCase  (	UpperCamelCase__   :   List[Any]       ,			UpperCamelCase__   :   int       ,			UpperCamelCase__   :   int					)      ->    Optional[Any]:

						"""simple docstring"""





						if split:
												__SCREAMING_SNAKE_CASE:				str           =   {split: jsonl_path}
						else:
												__SCREAMING_SNAKE_CASE:				Dict           =   '''train'''
												__SCREAMING_SNAKE_CASE:				Tuple           =   {'''train''': jsonl_path, '''test''': jsonl_path}
						__SCREAMING_SNAKE_CASE:				Tuple           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				str           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE:				int           =   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 		lowerCAmelCase  (	UpperCamelCase__   :   Dict					)      ->    Optional[int]:

						"""simple docstring"""





						return json.load(UpperCamelCase__					)






def 		lowerCAmelCase  (	UpperCamelCase__   :   Any					)      ->    Any:

						"""simple docstring"""





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





class 				a			:






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

										"""simple docstring"""




										with io.BytesIO() as buffer:
																JsonDatasetWriter(_lowerCAmelCase			, _lowerCAmelCase			, lines=_lowerCAmelCase					).write()
																buffer.seek(0					)
																__SCREAMING_SNAKE_CASE:				str           =   load_json_function(_lowerCAmelCase					)
										assert isinstance(_lowerCAmelCase			, _lowerCAmelCase					)
										assert isinstance(exported_content[0]			, _lowerCAmelCase					)
										assert len(_lowerCAmelCase					) == 10






				@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 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										with io.BytesIO() as buffer:
																JsonDatasetWriter(_lowerCAmelCase			, _lowerCAmelCase			, lines=_lowerCAmelCase			, orient=_lowerCAmelCase					).write()
																buffer.seek(0					)
																__SCREAMING_SNAKE_CASE:				Optional[int]           =   load_json(_lowerCAmelCase					)
										assert isinstance(_lowerCAmelCase			, _lowerCAmelCase					)
										if keys:
																if container is dict:
																						assert exported_content.keys() == keys
																else:
																						assert exported_content[0].keys() == keys
										else:
																assert not hasattr(_lowerCAmelCase			, '''keys'''					) and not hasattr(exported_content[0]			, '''keys'''					)
										if len_at:
																assert len(exported_content[len_at]					) == 10
										else:
																assert len(_lowerCAmelCase					) == 10






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

										"""simple docstring"""




										with io.BytesIO() as buffer:
																JsonDatasetWriter(_lowerCAmelCase			, _lowerCAmelCase			, lines=_lowerCAmelCase			, num_proc=2					).write()
																buffer.seek(0					)
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   load_json_function(_lowerCAmelCase					)
										assert isinstance(_lowerCAmelCase			, _lowerCAmelCase					)
										assert isinstance(exported_content[0]			, _lowerCAmelCase					)
										assert len(_lowerCAmelCase					) == 10






				@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 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										with io.BytesIO() as buffer:
																JsonDatasetWriter(_lowerCAmelCase			, _lowerCAmelCase			, lines=_lowerCAmelCase			, orient=_lowerCAmelCase			, num_proc=2					).write()
																buffer.seek(0					)
																__SCREAMING_SNAKE_CASE:				Dict           =   load_json(_lowerCAmelCase					)
										assert isinstance(_lowerCAmelCase			, _lowerCAmelCase					)
										if keys:
																if container is dict:
																						assert exported_content.keys() == keys
																else:
																						assert exported_content[0].keys() == keys
										else:
																assert not hasattr(_lowerCAmelCase			, '''keys'''					) and not hasattr(exported_content[0]			, '''keys'''					)
										if len_at:
																assert len(exported_content[len_at]					) == 10
										else:
																assert len(_lowerCAmelCase					) == 10






				def 			snake_case_    (  self			, _lowerCAmelCase					):

										"""simple docstring"""




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






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

										"""simple docstring"""




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

										with fsspec.open(_lowerCAmelCase			, '''rb'''			, compression='''infer'''					) as f:
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   f.read()
										with fsspec.open(_lowerCAmelCase			, '''rb'''			, compression='''infer'''					) as f:
																__SCREAMING_SNAKE_CASE:				Any           =   f.read()
										assert exported_content == original_content



