





import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
    import onnxruntime as ort

class 		SCREAMING_SNAKE_CASE__				(       _UpperCAmelCase						,      unittest.TestCase  ):
      A_						:				Tuple							=    'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'


      def    a (self       :      int				,			a__       :      List[Any]=0      ):





         """simple docstring"""



         __snake_case							       =   floats_tensor((1, 3, 128, 128)				,			rng=random.Random(a__      )      )
         __snake_case							       =   np.random.RandomState(a__      )
         __snake_case							       =   {
             '''prompt''': '''A painting of a squirrel eating a burger''',
             '''image''': image,
             '''generator''': generator,
             '''num_inference_steps''': 3,
             '''strength''': 0.7_5,
             '''guidance_scale''': 7.5,
             '''output_type''': '''numpy''',
         }
         return inputs


      def    a (self       :      List[Any]      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   self.get_dummy_inputs()
         __snake_case							       =   pipe(**a__      ).images
         __snake_case							       =   image[0, -3:, -3:, -1].flatten()

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.6_9_6_4_3, 0.5_8_4_8_4, 0.5_0_3_1_4, 0.5_8_7_6_0, 0.5_5_3_6_8, 0.5_9_6_4_3, 0.5_1_5_2_9, 0.4_1_2_1_7, 0.4_9_0_8_7]      )
         assert np.abs(image_slice - expected_slice      ).max() < 1E-1


      def    a (self       :      Dict      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   PNDMScheduler.from_config(pipe.scheduler.config				,			skip_prk_steps=a__      )
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   self.get_dummy_inputs()
         __snake_case							       =   pipe(**a__      ).images
         __snake_case							       =   image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.6_1_7_3_7, 0.5_4_6_4_2, 0.5_3_1_8_3, 0.5_4_4_6_5, 0.5_2_7_4_2, 0.6_0_5_2_5, 0.4_9_9_6_9, 0.4_0_6_5_5, 0.4_8_1_5_4]      )

         assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-1


      def    a (self       :      List[str]      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   LMSDiscreteScheduler.from_config(pipe.scheduler.config      )
         pipe.set_progress_bar_config(disable=a__      )

         # warmup pass to apply optimizations
         __snake_case							       =   pipe(**self.get_dummy_inputs()      )

         __snake_case							       =   self.get_dummy_inputs()
         __snake_case							       =   pipe(**a__      ).images
         __snake_case							       =   image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.5_2_7_6_1, 0.5_9_9_7_7, 0.4_9_0_3_3, 0.4_9_6_1_9, 0.5_4_2_8_2, 0.5_0_3_1_1, 0.4_7_6_0_0, 0.4_0_9_1_8, 0.4_5_2_0_3]      )

         assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-1


      def    a (self       :      Any      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   EulerDiscreteScheduler.from_config(pipe.scheduler.config      )
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   self.get_dummy_inputs()
         __snake_case							       =   pipe(**a__      ).images
         __snake_case							       =   image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]      )

         assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-1


      def    a (self       :      Dict      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config      )
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   self.get_dummy_inputs()
         __snake_case							       =   pipe(**a__      ).images
         __snake_case							       =   image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]      )

         assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-1


      def    a (self       :      List[str]      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config      )
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   self.get_dummy_inputs()
         __snake_case							       =   pipe(**a__      ).images
         __snake_case							       =   image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.6_5_3_3_1, 0.5_8_2_7_7, 0.4_8_2_0_4, 0.5_6_0_5_9, 0.5_3_6_6_5, 0.5_6_2_3_5, 0.5_0_9_6_9, 0.4_0_0_0_9, 0.4_6_5_5_2]      )

         assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-1






@nightly
@require_onnxruntime
@require_torch_gpu
class 		SCREAMING_SNAKE_CASE__				(       unittest.TestCase  ):


      @property
      def    a (self       :      List[str]      ):





         """simple docstring"""



         return (
             "CUDAExecutionProvider",
             {
                 "gpu_mem_limit": "15000000000",  # 15GB
                 "arena_extend_strategy": "kSameAsRequested",
             },
         )


      @property
      def    a (self       :      Optional[Any]      ):





         """simple docstring"""



         __snake_case							       =   ort.SessionOptions()
         __snake_case							       =   False
         return options


      def    a (self       :      Optional[Any]      ):





         """simple docstring"""



         __snake_case							       =   load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/img2img/sketch-mountains-input.jpg'''      )
         __snake_case							       =   init_image.resize((768, 512)      )
         # using the PNDM scheduler by default
         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(
             '''CompVis/stable-diffusion-v1-4'''				,			revision='''onnx'''				,			safety_checker=a__				,			feature_extractor=a__				,			provider=self.gpu_provider				,			sess_options=self.gpu_options				,			)
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   '''A fantasy landscape, trending on artstation'''

         __snake_case							       =   np.random.RandomState(0      )
         __snake_case							       =   pipe(
             prompt=a__				,			image=a__				,			strength=0.7_5				,			guidance_scale=7.5				,			num_inference_steps=10				,			generator=a__				,			output_type='''np'''				,			)
         __snake_case							       =   output.images
         __snake_case							       =   images[0, 255:258, 383:386, -1]

         assert images.shape == (1, 512, 768, 3)
         __snake_case							       =   np.array([0.4_9_0_9, 0.5_0_5_9, 0.5_3_7_2, 0.4_6_2_3, 0.4_8_7_6, 0.5_0_4_9, 0.4_8_2_0, 0.4_9_5_6, 0.5_0_1_9]      )
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

         assert np.abs(image_slice.flatten() - expected_slice      ).max() < 2E-2


      def    a (self       :      Dict      ):





         """simple docstring"""



         __snake_case							       =   load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/img2img/sketch-mountains-input.jpg'''      )
         __snake_case							       =   init_image.resize((768, 512)      )
         __snake_case							       =   LMSDiscreteScheduler.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''				,			subfolder='''scheduler'''				,			revision='''onnx'''      )
         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''				,			revision='''onnx'''				,			scheduler=a__				,			safety_checker=a__				,			feature_extractor=a__				,			provider=self.gpu_provider				,			sess_options=self.gpu_options				,			)
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   '''A fantasy landscape, trending on artstation'''

         __snake_case							       =   np.random.RandomState(0      )
         __snake_case							       =   pipe(
             prompt=a__				,			image=a__				,			strength=0.7_5				,			guidance_scale=7.5				,			num_inference_steps=20				,			generator=a__				,			output_type='''np'''				,			)
         __snake_case							       =   output.images
         __snake_case							       =   images[0, 255:258, 383:386, -1]

         assert images.shape == (1, 512, 768, 3)
         __snake_case							       =   np.array([0.8_0_4_3, 0.9_2_6, 0.9_5_8_1, 0.8_1_1_9, 0.8_9_5_4, 0.9_1_3, 0.7_2_0_9, 0.7_4_6_3, 0.7_4_3_1]      )
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

         assert np.abs(image_slice.flatten() - expected_slice      ).max() < 2E-2


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 			lowercase_				(			_lowerCamelCase							:       Optional[Any]		,	_lowerCamelCase							:       int):
    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 			lowercase_				(			_lowerCamelCase							:       List[Any]		,	_lowerCamelCase							:       Any		,	_lowerCamelCase							:       str):
    lowercase__ :   Optional[int]    					=				tmp_path / "cache"
    lowercase__ :   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():
        lowercase__ :   Union[str, Any]    					=				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 			lowercase_				(			_lowerCamelCase							:       List[str]		,	_lowerCamelCase							:       Dict		,	_lowerCamelCase							:       Dict):
    lowercase__ :   List[Any]    					=				tmp_path / "cache"
    lowercase__ :   Union[str, Any]    					=				{"col_1": "string", "col_2": "int64", "col_3": "float64"}
    lowercase__ :   List[Any]    					=				features.copy() if features else default_expected_features
    lowercase__ :   List[Any]    					=				(
        Features({feature: Value(_lowerCamelCase) for feature, dtype in features.items()}) if features is not None else None
    )
    lowercase__ :   Any    					=				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 			lowercase_				(			_lowerCamelCase							:       Any		,	_lowerCamelCase							:       Any		,	_lowerCamelCase							:       List[str]):
    lowercase__ :   Optional[Any]    					=				tmp_path / "cache"
    lowercase__ :   Tuple    					=				{"col_3": "float64", "col_1": "string", "col_2": "int64"}
    lowercase__ :   List[Any]    					=				features.copy() if features else default_expected_features
    lowercase__ :   int    					=				(
        Features({feature: Value(_lowerCamelCase) for feature, dtype in features.items()}) if features is not None else None
    )
    lowercase__ :   Any    					=				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 			lowercase_				(			_lowerCamelCase							:       List[Any]		,	_lowerCamelCase							:       Optional[int]):
    # jsonl_312_path features are {"col_3": "float64", "col_1": "string", "col_2": "int64"}
    lowercase__ :   Any    					=				{"col_2": "int64", "col_3": "float64", "col_1": "string"}
    lowercase__ :   str    					=				features.copy()
    lowercase__ :   str    					=				(
        Features({feature: Value(_lowerCamelCase) for feature, dtype in features.items()}) if features is not None else None
    )
    lowercase__ :   Optional[int]    					=				tmp_path / "cache"
    lowercase__ :   Any    					=				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 			lowercase_				(			_lowerCamelCase							:       Dict		,	_lowerCamelCase							:       Optional[int]		,	_lowerCamelCase							:       List[str]):
    lowercase__ :   Union[str, Any]    					=				tmp_path / "cache"
    lowercase__ :   List[Any]    					=				{"col_1": "string", "col_2": "int64", "col_3": "float64"}
    lowercase__ :   Union[str, Any]    					=				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 			lowercase_				(			_lowerCamelCase							:       Union[str, Any]		,	_lowerCamelCase							:       str		,	_lowerCamelCase							:       int):
    if issubclass(_lowerCamelCase		,	_lowerCamelCase):
        lowercase__ :   Tuple    					=				jsonl_path
    elif issubclass(_lowerCamelCase		,	_lowerCamelCase):
        lowercase__ :   str    					=				[jsonl_path]
    lowercase__ :   str    					=				tmp_path / "cache"
    lowercase__ :   Optional[Any]    					=				{"col_1": "string", "col_2": "int64", "col_3": "float64"}
    lowercase__ :   Tuple    					=				JsonDatasetReader(_lowerCamelCase		,	cache_dir=_lowerCamelCase).read()
    _check_json_dataset(_lowerCamelCase		,	_lowerCamelCase)

def 			lowercase_				(			_lowerCamelCase							:       List[Any]		,	_lowerCamelCase							:       Optional[Any]		,	_lowerCamelCase							:       Optional[int]=("train",)):
    assert isinstance(_lowerCamelCase		,	_lowerCamelCase)
    for split in splits:
        lowercase__ :   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 			lowercase_				(			_lowerCamelCase							:       Union[str, Any]		,	_lowerCamelCase							:       str		,	_lowerCamelCase							:       str):
    lowercase__ :   List[str]    					=				tmp_path / "cache"
    lowercase__ :   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():
        lowercase__ :   Optional[Any]    					=				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 			lowercase_				(			_lowerCamelCase							:       Any		,	_lowerCamelCase							:       List[str]		,	_lowerCamelCase							:       List[str]):
    lowercase__ :   str    					=				tmp_path / "cache"
    lowercase__ :   Tuple    					=				{"col_1": "string", "col_2": "int64", "col_3": "float64"}
    lowercase__ :   Tuple    					=				features.copy() if features else default_expected_features
    lowercase__ :   Union[str, Any]    					=				(
        Features({feature: Value(_lowerCamelCase) for feature, dtype in features.items()}) if features is not None else None
    )
    lowercase__ :   Tuple    					=				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 			lowercase_				(			_lowerCamelCase							:       str		,	_lowerCamelCase							:       Dict		,	_lowerCamelCase							:       Tuple):
    if split:
        lowercase__ :   Tuple    					=				{split: jsonl_path}
    else:
        lowercase__ :   Tuple    					=				"train"
        lowercase__ :   int    					=				{"train": jsonl_path, "test": jsonl_path}
    lowercase__ :   Dict    					=				tmp_path / "cache"
    lowercase__ :   Union[str, Any]    					=				{"col_1": "string", "col_2": "int64", "col_3": "float64"}
    lowercase__ :   Union[str, Any]    					=				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 			lowercase_				(			_lowerCamelCase							:       Union[str, Any]):
    return json.load(_lowerCamelCase)

def 			lowercase_				(			_lowerCamelCase							:       Optional[int]):
    return [json.loads(_lowerCamelCase) for line in buffer]



class   snake_case_       :






   @pytest.mark.parametrize("lines, load_json_function"	,						[(True, load_json_lines), (False, load_json)]  )
   def       __UpperCamelCase     (      self				:     List[Any]	,						lowercase_				:     Any	,						lowercase_				:     Optional[Any]	,						lowercase_				:     Dict  )					->						Optional[Any]:
       with io.BytesIO() as buffer:
           JsonDatasetWriter(lowercase_	,						lowercase_	,						lines=lowercase_  ).write()
           buffer.seek(0  )
           lowercase__ :   Optional[int]    					=				load_json_function(lowercase_  )
       assert isinstance(lowercase_	,						lowercase_  )
       assert isinstance(exported_content[0]	,						lowercase_  )
       assert len(lowercase_  ) == 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       __UpperCamelCase     (      self				:     str	,						lowercase_				:     int	,						lowercase_				:     str	,						lowercase_				:     Union[str, Any]	,						lowercase_				:     Dict	,						lowercase_				:     Tuple  )					->						List[str]:
       with io.BytesIO() as buffer:
           JsonDatasetWriter(lowercase_	,						lowercase_	,						lines=lowercase_	,						orient=lowercase_  ).write()
           buffer.seek(0  )
           lowercase__ :   str    					=				load_json(lowercase_  )
       assert isinstance(lowercase_	,						lowercase_  )
       if keys:
           if container is dict:
               assert exported_content.keys() == keys
           else:
               assert exported_content[0].keys() == keys
       else:
           assert not hasattr(lowercase_	,						"keys"  ) and not hasattr(exported_content[0]	,						"keys"  )
       if len_at:
           assert len(exported_content[len_at]  ) == 10
       else:
           assert len(lowercase_  ) == 10






   @pytest.mark.parametrize("lines, load_json_function"	,						[(True, load_json_lines), (False, load_json)]  )
   def       __UpperCamelCase     (      self				:     List[Any]	,						lowercase_				:     int	,						lowercase_				:     Optional[Any]	,						lowercase_				:     Dict  )					->						Optional[int]:
       with io.BytesIO() as buffer:
           JsonDatasetWriter(lowercase_	,						lowercase_	,						lines=lowercase_	,						num_proc=2  ).write()
           buffer.seek(0  )
           lowercase__ :   str    					=				load_json_function(lowercase_  )
       assert isinstance(lowercase_	,						lowercase_  )
       assert isinstance(exported_content[0]	,						lowercase_  )
       assert len(lowercase_  ) == 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       __UpperCamelCase     (      self				:     Union[str, Any]	,						lowercase_				:     Dict	,						lowercase_				:     Dict	,						lowercase_				:     Optional[int]	,						lowercase_				:     Optional[Any]	,						lowercase_				:     Dict  )					->						Any:
       with io.BytesIO() as buffer:
           JsonDatasetWriter(lowercase_	,						lowercase_	,						lines=lowercase_	,						orient=lowercase_	,						num_proc=2  ).write()
           buffer.seek(0  )
           lowercase__ :   Optional[Any]    					=				load_json(lowercase_  )
       assert isinstance(lowercase_	,						lowercase_  )
       if keys:
           if container is dict:
               assert exported_content.keys() == keys
           else:
               assert exported_content[0].keys() == keys
       else:
           assert not hasattr(lowercase_	,						"keys"  ) and not hasattr(exported_content[0]	,						"keys"  )
       if len_at:
           assert len(exported_content[len_at]  ) == 10
       else:
           assert len(lowercase_  ) == 10






   def       __UpperCamelCase     (      self				:     Dict	,						lowercase_				:     List[str]  )					->						str:
       with pytest.raises(lowercase_  ):
           with io.BytesIO() as buffer:
               JsonDatasetWriter(lowercase_	,						lowercase_	,						num_proc=0  )







   @pytest.mark.parametrize("compression, extension"	,						[("gzip", "gz"), ("bz2", "bz2"), ("xz", "xz")]  )
   def       __UpperCamelCase     (      self				:     List[Any]	,						lowercase_				:     Tuple	,						lowercase_				:     List[str]	,						lowercase_				:     Optional[Any]	,						lowercase_				:     str	,						lowercase_				:     List[Any]  )					->						Any:
       lowercase__ :   Dict    					=				tmp_path_factory.mktemp("data"  ) / F'''test.json.{extension}'''
       lowercase__ :   Optional[int]    					=				str(shared_datadir / F'''test_file.json.{extension}'''  )
       JsonDatasetWriter(lowercase_	,						lowercase_	,						compression=lowercase_  ).write()

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



