


"""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     __lowerCAmelCase		(		lowercase    :    Union[str, Any]		,				lowercase    :    Optional[int]					)      ->		Union[str, Any]:


    """simple docstring"""





    assert isinstance(lowercase		,				lowercase					)
    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		(		lowercase    :    Dict		,				lowercase    :    Any		,				lowercase    :    List[Any]					)      ->		List[Any]:


    """simple docstring"""





    snake_case       : Union[str, Any]    =       tmp_path / "cache"
    snake_case       : Optional[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():
        snake_case       : Optional[int]    =       JsonDatasetReader(lowercase		,				cache_dir=lowercase		,				keep_in_memory=lowercase					).read()
    _check_json_dataset(lowercase		,				lowercase					)







@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		(		lowercase    :    int		,				lowercase    :    int		,				lowercase    :    Union[str, Any]					)      ->		Dict:


    """simple docstring"""





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







@pytest.mark.parametrize(
    "features"		,				[
        None,
        {"col_3": "float64", "col_1": "string", "col_2": "int64"},
    ]		,				)
def     __lowerCAmelCase		(		lowercase    :    Tuple		,				lowercase    :    int		,				lowercase    :    Dict					)      ->		List[Any]:


    """simple docstring"""





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


    """simple docstring"""





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


    """simple docstring"""





    snake_case       : Optional[Any]    =       tmp_path / "cache"
    snake_case       : List[Any]    =       {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    snake_case       : Dict    =       JsonDatasetReader(lowercase		,				cache_dir=lowercase		,				split=lowercase					).read()
    _check_json_dataset(lowercase		,				lowercase					)
    assert dataset.split == split if split else "train"







@pytest.mark.parametrize("path_type"		,				[str, list]					)
def     __lowerCAmelCase		(		lowercase    :    List[Any]		,				lowercase    :    List[str]		,				lowercase    :    Optional[Any]					)      ->		Any:


    """simple docstring"""





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







def     __lowerCAmelCase		(		lowercase    :    int		,				lowercase    :    Any		,				lowercase    :    List[Any]=("train",)					)      ->		Union[str, Any]:


    """simple docstring"""





    assert isinstance(lowercase		,				lowercase					)
    for split in splits:
        snake_case       : Optional[Any]    =       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		(		lowercase    :    Optional[int]		,				lowercase    :    Dict		,				lowercase    :    Union[str, Any]					)      ->		Optional[int]:


    """simple docstring"""





    snake_case       : Any    =       tmp_path / "cache"
    snake_case       : 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():
        snake_case       : Tuple    =       JsonDatasetReader({"train": jsonl_path}		,				cache_dir=lowercase		,				keep_in_memory=lowercase					).read()
    _check_json_datasetdict(lowercase		,				lowercase					)







@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		(		lowercase    :    List[Any]		,				lowercase    :    List[Any]		,				lowercase    :    Dict					)      ->		Tuple:


    """simple docstring"""





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







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


    """simple docstring"""





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







def     __lowerCAmelCase		(		lowercase    :    List[str]					)      ->		Dict:


    """simple docstring"""





    return json.load(lowercase					)







def     __lowerCAmelCase		(		lowercase    :    List[str]					)      ->		int:


    """simple docstring"""





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







class _lowerCAmelCase							:







   @pytest.mark.parametrize("lines, load_json_function"		,      [(True, load_json_lines), (False, load_json)]		)
   def 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		)      ->						Dict:

       '''simple docstring'''




       with io.BytesIO() as buffer:
           JsonDatasetWriter(UpperCamelCase__		,      UpperCamelCase__		,      lines=UpperCamelCase__		).write()
           buffer.seek(0		)
           snake_case       : Tuple    =       load_json_function(UpperCamelCase__		)
       assert isinstance(UpperCamelCase__		,      UpperCamelCase__		)
       assert isinstance(exported_content[0]		,      UpperCamelCase__		)
       assert len(UpperCamelCase__		) == 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 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		)      ->						List[Any]:

       '''simple docstring'''




       with io.BytesIO() as buffer:
           JsonDatasetWriter(UpperCamelCase__		,      UpperCamelCase__		,      lines=UpperCamelCase__		,      orient=UpperCamelCase__		).write()
           buffer.seek(0		)
           snake_case       : str    =       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]		) == 10
       else:
           assert len(UpperCamelCase__		) == 10







   @pytest.mark.parametrize("lines, load_json_function"		,      [(True, load_json_lines), (False, load_json)]		)
   def 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		)      ->						List[str]:

       '''simple docstring'''




       with io.BytesIO() as buffer:
           JsonDatasetWriter(UpperCamelCase__		,      UpperCamelCase__		,      lines=UpperCamelCase__		,      num_proc=2		).write()
           buffer.seek(0		)
           snake_case       : str    =       load_json_function(UpperCamelCase__		)
       assert isinstance(UpperCamelCase__		,      UpperCamelCase__		)
       assert isinstance(exported_content[0]		,      UpperCamelCase__		)
       assert len(UpperCamelCase__		) == 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 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		)      ->						Optional[int]:

       '''simple docstring'''




       with io.BytesIO() as buffer:
           JsonDatasetWriter(UpperCamelCase__		,      UpperCamelCase__		,      lines=UpperCamelCase__		,      orient=UpperCamelCase__		,      num_proc=2		).write()
           buffer.seek(0		)
           snake_case       : List[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]		) == 10
       else:
           assert len(UpperCamelCase__		) == 10







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		)      ->						str:

       '''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 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		,      UpperCamelCase__		)      ->						List[str]:

       '''simple docstring'''




       snake_case       : Optional[int]    =       tmp_path_factory.mktemp("data"		) / F'test.json.{extension}'
       snake_case       : List[str]    =       str(shared_datadir / F'test_file.json.{extension}'		)
       JsonDatasetWriter(UpperCamelCase__		,      UpperCamelCase__		,      compression=UpperCamelCase__		).write()

       with fsspec.open(UpperCamelCase__		,      "rb"		,      compression="infer"		) as f:
           snake_case       : List[Any]    =       f.read()
       with fsspec.open(UpperCamelCase__		,      "rb"		,      compression="infer"		) as f:
           snake_case       : List[Any]    =       f.read()
       assert exported_content == original_content









"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


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

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __snake_case       =       [
      """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

  __snake_case       =       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





