



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_(lowerCamelCase_    ,   lowerCamelCase_    ) ->       Any:
						assert isinstance(lowerCamelCase_    ,   lowerCamelCase_    )
						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_(lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_    ) ->       Optional[int]:
						UpperCAmelCase         =     tmp_path / "cache"
						UpperCAmelCase         =     {"col_1": "string", "col_2": "int64", "col_3": "float64"}
						with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
												UpperCAmelCase         =     JsonDatasetReader(lowerCamelCase_    ,   cache_dir=lowerCamelCase_    ,   keep_in_memory=lowerCamelCase_    ).read()
						_check_json_dataset(lowerCamelCase_    ,   lowerCamelCase_    )

@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_(lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_    ) ->       Tuple:
						UpperCAmelCase         =     tmp_path / "cache"
						UpperCAmelCase         =     {"col_1": "string", "col_2": "int64", "col_3": "float64"}
						UpperCAmelCase         =     features.copy() if features else default_expected_features
						UpperCAmelCase         =     (
						    Features({feature: Value(lowerCamelCase_    ) for feature, dtype in features.items()}    ) if features is not None else None
						)
						UpperCAmelCase         =     JsonDatasetReader(lowerCamelCase_    ,   features=lowerCamelCase_    ,   cache_dir=lowerCamelCase_    ).read()
						_check_json_dataset(lowerCamelCase_    ,   lowerCamelCase_    )

@pytest.mark.parametrize(
    "features"    ,   [
        None,
        {"col_3": "float64", "col_1": "string", "col_2": "int64"},
    ]    ,   )
def    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_    ) ->       Tuple:
						UpperCAmelCase         =     tmp_path / "cache"
						UpperCAmelCase         =     {"col_3": "float64", "col_1": "string", "col_2": "int64"}
						UpperCAmelCase         =     features.copy() if features else default_expected_features
						UpperCAmelCase         =     (
						    Features({feature: Value(lowerCamelCase_    ) for feature, dtype in features.items()}    ) if features is not None else None
						)
						UpperCAmelCase         =     JsonDatasetReader(lowerCamelCase_    ,   features=lowerCamelCase_    ,   cache_dir=lowerCamelCase_    ).read()
						assert isinstance(lowerCamelCase_    ,   lowerCamelCase_    )
						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_(lowerCamelCase_    ,   lowerCamelCase_    ) ->       Any:
						# jsonl_312_path features are {"col_3": "float64", "col_1": "string", "col_2": "int64"}
						UpperCAmelCase         =     {"col_2": "int64", "col_3": "float64", "col_1": "string"}
						UpperCAmelCase         =     features.copy()
						UpperCAmelCase         =     (
						    Features({feature: Value(lowerCamelCase_    ) for feature, dtype in features.items()}    ) if features is not None else None
						)
						UpperCAmelCase         =     tmp_path / "cache"
						UpperCAmelCase         =     JsonDatasetReader(lowerCamelCase_    ,   features=lowerCamelCase_    ,   cache_dir=lowerCamelCase_    ).read()
						assert isinstance(lowerCamelCase_    ,   lowerCamelCase_    )
						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_(lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_    ) ->       List[str]:
						UpperCAmelCase         =     tmp_path / "cache"
						UpperCAmelCase         =     {"col_1": "string", "col_2": "int64", "col_3": "float64"}
						UpperCAmelCase         =     JsonDatasetReader(lowerCamelCase_    ,   cache_dir=lowerCamelCase_    ,   split=lowerCamelCase_    ).read()
						_check_json_dataset(lowerCamelCase_    ,   lowerCamelCase_    )
						assert dataset.split == split if split else "train"

@pytest.mark.parametrize("path_type"    ,   [str, list]    )
def    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_    ) ->       List[str]:
						if issubclass(lowerCamelCase_    ,   lowerCamelCase_    ):
												UpperCAmelCase         =     jsonl_path
						elif issubclass(lowerCamelCase_    ,   lowerCamelCase_    ):
												UpperCAmelCase         =     [jsonl_path]
						UpperCAmelCase         =     tmp_path / "cache"
						UpperCAmelCase         =     {"col_1": "string", "col_2": "int64", "col_3": "float64"}
						UpperCAmelCase         =     JsonDatasetReader(lowerCamelCase_    ,   cache_dir=lowerCamelCase_    ).read()
						_check_json_dataset(lowerCamelCase_    ,   lowerCamelCase_    )

def    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_=("train",)    ) ->       List[Any]:
						assert isinstance(lowerCamelCase_    ,   lowerCamelCase_    )
						for split in splits:
												UpperCAmelCase         =     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_(lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_    ) ->       Union[str, Any]:
						UpperCAmelCase         =     tmp_path / "cache"
						UpperCAmelCase         =     {"col_1": "string", "col_2": "int64", "col_3": "float64"}
						with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
												UpperCAmelCase         =     JsonDatasetReader({"train": jsonl_path}    ,   cache_dir=lowerCamelCase_    ,   keep_in_memory=lowerCamelCase_    ).read()
						_check_json_datasetdict(lowerCamelCase_    ,   lowerCamelCase_    )

@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_(lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_    ) ->       int:
						UpperCAmelCase         =     tmp_path / "cache"
						UpperCAmelCase         =     {"col_1": "string", "col_2": "int64", "col_3": "float64"}
						UpperCAmelCase         =     features.copy() if features else default_expected_features
						UpperCAmelCase         =     (
						    Features({feature: Value(lowerCamelCase_    ) for feature, dtype in features.items()}    ) if features is not None else None
						)
						UpperCAmelCase         =     JsonDatasetReader({"train": jsonl_path}    ,   features=lowerCamelCase_    ,   cache_dir=lowerCamelCase_    ).read()
						_check_json_datasetdict(lowerCamelCase_    ,   lowerCamelCase_    )

@pytest.mark.parametrize("split"    ,   [None, NamedSplit("train"    ), "train", "test"]    )
def    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_    ) ->       Dict:
						if split:
												UpperCAmelCase         =     {split: jsonl_path}
						else:
												UpperCAmelCase         =     "train"
												UpperCAmelCase         =     {"train": jsonl_path, "test": jsonl_path}
						UpperCAmelCase         =     tmp_path / "cache"
						UpperCAmelCase         =     {"col_1": "string", "col_2": "int64", "col_3": "float64"}
						UpperCAmelCase         =     JsonDatasetReader(lowerCamelCase_    ,   cache_dir=lowerCamelCase_    ).read()
						_check_json_datasetdict(lowerCamelCase_    ,   lowerCamelCase_    ,   splits=list(path.keys()    )    )
						assert all(dataset[split].split == split for split in path.keys()    )

def    lowerCamelCase_(lowerCamelCase_    ) ->       Dict:
						return json.load(lowerCamelCase_    )

def    lowerCamelCase_(lowerCamelCase_    ) ->       str:
						return [json.loads(lowerCamelCase_    ) for line in buffer]





class  __magic_name__   :






		@pytest.mark.parametrize("lines, load_json_function"				,		[(True, load_json_lines), (False, load_json)]      )
		def SCREAMING_SNAKE_CASE_			(						self      :     List[Any]				,		UpperCamelCase__      :     int				,		UpperCamelCase__      :     Optional[int]				,		UpperCamelCase__      :     Any      )				->					List[str]:
								'''simple docstring'''







								with io.BytesIO() as buffer:
														JsonDatasetWriter(UpperCamelCase__				,		UpperCamelCase__				,		lines=UpperCamelCase__      ).write()
														buffer.seek(0      )
														UpperCAmelCase         =     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 SCREAMING_SNAKE_CASE_			(						self      :     Optional[Any]				,		UpperCamelCase__      :     Dict				,		UpperCamelCase__      :     Dict				,		UpperCamelCase__      :     Dict				,		UpperCamelCase__      :     Optional[int]				,		UpperCamelCase__      :     List[Any]      )				->					Any:
								'''simple docstring'''







								with io.BytesIO() as buffer:
														JsonDatasetWriter(UpperCamelCase__				,		UpperCamelCase__				,		lines=UpperCamelCase__				,		orient=UpperCamelCase__      ).write()
														buffer.seek(0      )
														UpperCAmelCase         =     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 SCREAMING_SNAKE_CASE_			(						self      :     int				,		UpperCamelCase__      :     Optional[Any]				,		UpperCamelCase__      :     Union[str, Any]				,		UpperCamelCase__      :     List[str]      )				->					Optional[Any]:
								'''simple docstring'''







								with io.BytesIO() as buffer:
														JsonDatasetWriter(UpperCamelCase__				,		UpperCamelCase__				,		lines=UpperCamelCase__				,		num_proc=2      ).write()
														buffer.seek(0      )
														UpperCAmelCase         =     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 SCREAMING_SNAKE_CASE_			(						self      :     Optional[Any]				,		UpperCamelCase__      :     Dict				,		UpperCamelCase__      :     Any				,		UpperCamelCase__      :     Any				,		UpperCamelCase__      :     str				,		UpperCamelCase__      :     str      )				->					Optional[Any]:
								'''simple docstring'''







								with io.BytesIO() as buffer:
														JsonDatasetWriter(UpperCamelCase__				,		UpperCamelCase__				,		lines=UpperCamelCase__				,		orient=UpperCamelCase__				,		num_proc=2      ).write()
														buffer.seek(0      )
														UpperCAmelCase         =     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 SCREAMING_SNAKE_CASE_			(						self      :     Union[str, Any]				,		UpperCamelCase__      :     Union[str, Any]      )				->					Optional[int]:
								'''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 SCREAMING_SNAKE_CASE_			(						self      :     Union[str, Any]				,		UpperCamelCase__      :     str				,		UpperCamelCase__      :     Any				,		UpperCamelCase__      :     int				,		UpperCamelCase__      :     Optional[Any]				,		UpperCamelCase__      :     List[Any]      )				->					Union[str, Any]:
								'''simple docstring'''







								UpperCAmelCase         =     tmp_path_factory.mktemp("data"      ) / F'test.json.{extension}'
								UpperCAmelCase         =     str(shared_datadir / F'test_file.json.{extension}'      )
								JsonDatasetWriter(UpperCamelCase__				,		UpperCamelCase__				,		compression=UpperCamelCase__      ).write()

								with fsspec.open(UpperCamelCase__				,		"rb"				,		compression="infer"      ) as f:
														UpperCAmelCase         =     f.read()
								with fsspec.open(UpperCamelCase__				,		"rb"				,		compression="infer"      ) as f:
														UpperCAmelCase         =     f.read()
								assert exported_content == original_content








import unittest

import torch
from torch import nn

from diffusers.models.activations import get_activation





class  __magic_name__   (							unittest.TestCase						):






		def SCREAMING_SNAKE_CASE_			(						self      :     int      )				->					str:
								'''simple docstring'''







								UpperCAmelCase         =     get_activation("swish"      )

								self.assertIsInstance(UpperCamelCase__				,		nn.SiLU      )

								self.assertEqual(act(torch.tensor(-1_00				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertNotEqual(act(torch.tensor(-1				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(0				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(20				,		dtype=torch.floataa      )      ).item()				,		20      )






		def SCREAMING_SNAKE_CASE_			(						self      :     Tuple      )				->					List[Any]:
								'''simple docstring'''







								UpperCAmelCase         =     get_activation("silu"      )

								self.assertIsInstance(UpperCamelCase__				,		nn.SiLU      )

								self.assertEqual(act(torch.tensor(-1_00				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertNotEqual(act(torch.tensor(-1				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(0				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(20				,		dtype=torch.floataa      )      ).item()				,		20      )






		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[Any]      )				->					int:
								'''simple docstring'''







								UpperCAmelCase         =     get_activation("mish"      )

								self.assertIsInstance(UpperCamelCase__				,		nn.Mish      )

								self.assertEqual(act(torch.tensor(-2_00				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertNotEqual(act(torch.tensor(-1				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(0				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(20				,		dtype=torch.floataa      )      ).item()				,		20      )







		def SCREAMING_SNAKE_CASE_			(						self      :     str      )				->					List[str]:
								'''simple docstring'''







								UpperCAmelCase         =     get_activation("gelu"      )

								self.assertIsInstance(UpperCamelCase__				,		nn.GELU      )

								self.assertEqual(act(torch.tensor(-1_00				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertNotEqual(act(torch.tensor(-1				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(0				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(20				,		dtype=torch.floataa      )      ).item()				,		20      )



