





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 _a (  SCREAMING_SNAKE_CASE__      :  List[str]			,				SCREAMING_SNAKE_CASE__      :  Tuple       )	->    List[Any]:

							'''simple docstring'''






							assert isinstance(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       )
							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 _a (  SCREAMING_SNAKE_CASE__      :  Optional[int]			,				SCREAMING_SNAKE_CASE__      :  Dict			,				SCREAMING_SNAKE_CASE__      :  Dict       )	->    Optional[int]:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      Optional[int]     =						tmp_path / "cache"
							SCREAMING_SNAKE_CASE__       :      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__       :      Optional[int]     =						JsonDatasetReader(SCREAMING_SNAKE_CASE__			,				cache_dir=SCREAMING_SNAKE_CASE__			,				keep_in_memory=SCREAMING_SNAKE_CASE__       ).read()
							_check_json_dataset(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       )




@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 _a (  SCREAMING_SNAKE_CASE__      :  int			,				SCREAMING_SNAKE_CASE__      :  Tuple			,				SCREAMING_SNAKE_CASE__      :  List[Any]       )	->    Union[str, Any]:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      int     =						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(SCREAMING_SNAKE_CASE__       ) for feature, dtype in features.items()}       ) if features is not None else None
							)
							SCREAMING_SNAKE_CASE__       :      Optional[int]     =						JsonDatasetReader(SCREAMING_SNAKE_CASE__			,				features=SCREAMING_SNAKE_CASE__			,				cache_dir=SCREAMING_SNAKE_CASE__       ).read()
							_check_json_dataset(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       )




@pytest.mark.parametrize(
    "features"			,				[
        None,
        {"col_3": "float64", "col_1": "string", "col_2": "int64"},
    ]			,				)
def _a (  SCREAMING_SNAKE_CASE__      :  Dict			,				SCREAMING_SNAKE_CASE__      :  str			,				SCREAMING_SNAKE_CASE__      :  Tuple       )	->    Optional[int]:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      Optional[int]     =						tmp_path / "cache"
							SCREAMING_SNAKE_CASE__       :      str     =						{"col_3": "float64", "col_1": "string", "col_2": "int64"}
							SCREAMING_SNAKE_CASE__       :      Any     =						features.copy() if features else default_expected_features
							SCREAMING_SNAKE_CASE__       :      Dict     =						(
							    Features({feature: Value(SCREAMING_SNAKE_CASE__       ) for feature, dtype in features.items()}       ) if features is not None else None
							)
							SCREAMING_SNAKE_CASE__       :      Any     =						JsonDatasetReader(SCREAMING_SNAKE_CASE__			,				features=SCREAMING_SNAKE_CASE__			,				cache_dir=SCREAMING_SNAKE_CASE__       ).read()
							assert isinstance(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       )
							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 _a (  SCREAMING_SNAKE_CASE__      :  Any			,				SCREAMING_SNAKE_CASE__      :  int       )	->    List[Any]:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      Dict     =						{"col_2": "int64", "col_3": "float64", "col_1": "string"}
							SCREAMING_SNAKE_CASE__       :      Any     =						features.copy()
							SCREAMING_SNAKE_CASE__       :      str     =						(
							    Features({feature: Value(SCREAMING_SNAKE_CASE__       ) 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__       :      int     =						JsonDatasetReader(SCREAMING_SNAKE_CASE__			,				features=SCREAMING_SNAKE_CASE__			,				cache_dir=SCREAMING_SNAKE_CASE__       ).read()
							assert isinstance(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       )
							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 _a (  SCREAMING_SNAKE_CASE__      :  Any			,				SCREAMING_SNAKE_CASE__      :  str			,				SCREAMING_SNAKE_CASE__      :  Union[str, Any]       )	->    str:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      Tuple     =						tmp_path / "cache"
							SCREAMING_SNAKE_CASE__       :      Any     =						{"col_1": "string", "col_2": "int64", "col_3": "float64"}
							SCREAMING_SNAKE_CASE__       :      Dict     =						JsonDatasetReader(SCREAMING_SNAKE_CASE__			,				cache_dir=SCREAMING_SNAKE_CASE__			,				split=SCREAMING_SNAKE_CASE__       ).read()
							_check_json_dataset(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       )
							assert dataset.split == split if split else "train"




@pytest.mark.parametrize("path_type"			,				[str, list]       )
def _a (  SCREAMING_SNAKE_CASE__      :  Dict			,				SCREAMING_SNAKE_CASE__      :  Optional[Any]			,				SCREAMING_SNAKE_CASE__      :  int       )	->    Optional[int]:

							'''simple docstring'''






							if issubclass(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       ):
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						jsonl_path
							elif issubclass(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       ):
														SCREAMING_SNAKE_CASE__       :      int     =						[jsonl_path]
							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__       :      Any     =						JsonDatasetReader(SCREAMING_SNAKE_CASE__			,				cache_dir=SCREAMING_SNAKE_CASE__       ).read()
							_check_json_dataset(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       )




def _a (  SCREAMING_SNAKE_CASE__      :  Any			,				SCREAMING_SNAKE_CASE__      :  Optional[Any]			,				SCREAMING_SNAKE_CASE__      :  Optional[Any]=("train",)       )	->    int:

							'''simple docstring'''






							assert isinstance(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       )
							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 _a (  SCREAMING_SNAKE_CASE__      :  List[Any]			,				SCREAMING_SNAKE_CASE__      :  Any			,				SCREAMING_SNAKE_CASE__      :  Optional[int]       )	->    str:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      int     =						tmp_path / "cache"
							SCREAMING_SNAKE_CASE__       :      List[Any]     =						{"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__       :      Optional[Any]     =						JsonDatasetReader({"train": jsonl_path}			,				cache_dir=SCREAMING_SNAKE_CASE__			,				keep_in_memory=SCREAMING_SNAKE_CASE__       ).read()
							_check_json_datasetdict(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       )




@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 _a (  SCREAMING_SNAKE_CASE__      :  Optional[Any]			,				SCREAMING_SNAKE_CASE__      :  int			,				SCREAMING_SNAKE_CASE__      :  Optional[int]       )	->    Optional[int]:

							'''simple docstring'''






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




@pytest.mark.parametrize("split"			,				[None, NamedSplit("train"       ), "train", "test"]       )
def _a (  SCREAMING_SNAKE_CASE__      :  Optional[Any]			,				SCREAMING_SNAKE_CASE__      :  Dict			,				SCREAMING_SNAKE_CASE__      :  List[Any]       )	->    int:

							'''simple docstring'''






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




def _a (  SCREAMING_SNAKE_CASE__      :  List[Any]       )	->    int:

							'''simple docstring'''






							return json.load(SCREAMING_SNAKE_CASE__       )




def _a (  SCREAMING_SNAKE_CASE__      :  Optional[Any]       )	->    Union[str, Any]:

							'''simple docstring'''






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







class 						lowerCamelCase       :







							"""simple docstring"""




							@pytest.mark.parametrize("lines, load_json_function",			[(True, load_json_lines), (False, load_json)]						)
							def  A_     ( self     :					Optional[int],			_UpperCAmelCase     :					Union[str, Any],			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					Any						)				->			Optional[int]:





														"""simple docstring"""
														with io.BytesIO() as buffer:
																					JsonDatasetWriter(_UpperCAmelCase,			_UpperCAmelCase,			lines=_UpperCAmelCase						).write()
																					buffer.seek(0						)
																					SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						load_json_function(_UpperCAmelCase						)
														assert isinstance(_UpperCAmelCase,			_UpperCAmelCase						)
														assert isinstance(exported_content[0],			_UpperCAmelCase						)
														assert len(_UpperCAmelCase						) == 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  A_     ( self     :					Tuple,			_UpperCAmelCase     :					int,			_UpperCAmelCase     :					Union[str, Any],			_UpperCAmelCase     :					List[str],			_UpperCAmelCase     :					str,			_UpperCAmelCase     :					Optional[int]						)				->			Optional[int]:





														"""simple docstring"""
														with io.BytesIO() as buffer:
																					JsonDatasetWriter(_UpperCAmelCase,			_UpperCAmelCase,			lines=_UpperCAmelCase,			orient=_UpperCAmelCase						).write()
																					buffer.seek(0						)
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						load_json(_UpperCAmelCase						)
														assert isinstance(_UpperCAmelCase,			_UpperCAmelCase						)
														if keys:
																					if container is dict:
																												assert exported_content.keys() == keys
																					else:
																												assert exported_content[0].keys() == keys
														else:
																					assert not hasattr(_UpperCAmelCase,			"keys"						) and not hasattr(exported_content[0],			"keys"						)
														if len_at:
																					assert len(exported_content[len_at]						) == 1_0
														else:
																					assert len(_UpperCAmelCase						) == 1_0






							@pytest.mark.parametrize("lines, load_json_function",			[(True, load_json_lines), (False, load_json)]						)
							def  A_     ( self     :					List[str],			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					Dict,			_UpperCAmelCase     :					Optional[Any]						)				->			Optional[int]:





														"""simple docstring"""
														with io.BytesIO() as buffer:
																					JsonDatasetWriter(_UpperCAmelCase,			_UpperCAmelCase,			lines=_UpperCAmelCase,			num_proc=2						).write()
																					buffer.seek(0						)
																					SCREAMING_SNAKE_CASE__       :      str     =						load_json_function(_UpperCAmelCase						)
														assert isinstance(_UpperCAmelCase,			_UpperCAmelCase						)
														assert isinstance(exported_content[0],			_UpperCAmelCase						)
														assert len(_UpperCAmelCase						) == 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  A_     ( self     :					Optional[int],			_UpperCAmelCase     :					Tuple,			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					Union[str, Any],			_UpperCAmelCase     :					Tuple,			_UpperCAmelCase     :					Tuple						)				->			Optional[int]:





														"""simple docstring"""
														with io.BytesIO() as buffer:
																					JsonDatasetWriter(_UpperCAmelCase,			_UpperCAmelCase,			lines=_UpperCAmelCase,			orient=_UpperCAmelCase,			num_proc=2						).write()
																					buffer.seek(0						)
																					SCREAMING_SNAKE_CASE__       :      Optional[int]     =						load_json(_UpperCAmelCase						)
														assert isinstance(_UpperCAmelCase,			_UpperCAmelCase						)
														if keys:
																					if container is dict:
																												assert exported_content.keys() == keys
																					else:
																												assert exported_content[0].keys() == keys
														else:
																					assert not hasattr(_UpperCAmelCase,			"keys"						) and not hasattr(exported_content[0],			"keys"						)
														if len_at:
																					assert len(exported_content[len_at]						) == 1_0
														else:
																					assert len(_UpperCAmelCase						) == 1_0






							def  A_     ( self     :					Optional[Any],			_UpperCAmelCase     :					Optional[int]						)				->			Dict:





														"""simple docstring"""
														with pytest.raises(_UpperCAmelCase						):
																					with io.BytesIO() as buffer:
																												JsonDatasetWriter(_UpperCAmelCase,			_UpperCAmelCase,			num_proc=0						)






							@pytest.mark.parametrize("compression, extension",			[("gzip", "gz"), ("bz2", "bz2"), ("xz", "xz")]						)
							def  A_     ( self     :					Dict,			_UpperCAmelCase     :					List[str],			_UpperCAmelCase     :					Tuple,			_UpperCAmelCase     :					Dict,			_UpperCAmelCase     :					str,			_UpperCAmelCase     :					str						)				->			Optional[Any]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						tmp_path_factory.mktemp("data"						) / F'''test.json.{extension}'''
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						str(shared_datadir / F'''test_file.json.{extension}'''						)
														JsonDatasetWriter(_UpperCAmelCase,			_UpperCAmelCase,			compression=_UpperCAmelCase						).write()

														with fsspec.open(_UpperCAmelCase,			"rb",			compression="infer"						) as f:
																					SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						f.read()
														with fsspec.open(_UpperCAmelCase,			"rb",			compression="infer"						) as f:
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						f.read()
														assert exported_content == original_content













from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_lowerCamelCase :    Union[str, Any]  							= {
    '''configuration_bigbird_pegasus''': [
        '''BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''BigBirdPegasusConfig''',
        '''BigBirdPegasusOnnxConfig''',
    ],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowerCamelCase :    int  							= [
		    '''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 :    Optional[Any]  							= _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






