



import argparse
import os

import torch

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)


UpperCAmelCase							 =					{
    """sample_size""": 32,
    """in_channels""": 3,
    """out_channels""": 3,
    """layers_per_block""": 2,
    """num_class_embeds""": 1_000,
    """block_out_channels""": [32, 64],
    """attention_head_dim""": 8,
    """down_block_types""": [
        """ResnetDownsampleBlock2D""",
        """AttnDownBlock2D""",
    ],
    """up_block_types""": [
        """AttnUpBlock2D""",
        """ResnetUpsampleBlock2D""",
    ],
    """resnet_time_scale_shift""": """scale_shift""",
    """upsample_type""": """resnet""",
    """downsample_type""": """resnet""",
}

UpperCAmelCase							 =					{
    """sample_size""": 64,
    """in_channels""": 3,
    """out_channels""": 3,
    """layers_per_block""": 3,
    """num_class_embeds""": 1_000,
    """block_out_channels""": [192, 192 * 2, 192 * 3, 192 * 4],
    """attention_head_dim""": 64,
    """down_block_types""": [
        """ResnetDownsampleBlock2D""",
        """AttnDownBlock2D""",
        """AttnDownBlock2D""",
        """AttnDownBlock2D""",
    ],
    """up_block_types""": [
        """AttnUpBlock2D""",
        """AttnUpBlock2D""",
        """AttnUpBlock2D""",
        """ResnetUpsampleBlock2D""",
    ],
    """resnet_time_scale_shift""": """scale_shift""",
    """upsample_type""": """resnet""",
    """downsample_type""": """resnet""",
}

UpperCAmelCase							 =					{
    """sample_size""": 256,
    """in_channels""": 3,
    """out_channels""": 3,
    """layers_per_block""": 2,
    """num_class_embeds""": None,
    """block_out_channels""": [256, 256, 256 * 2, 256 * 2, 256 * 4, 256 * 4],
    """attention_head_dim""": 64,
    """down_block_types""": [
        """ResnetDownsampleBlock2D""",
        """ResnetDownsampleBlock2D""",
        """ResnetDownsampleBlock2D""",
        """AttnDownBlock2D""",
        """AttnDownBlock2D""",
        """AttnDownBlock2D""",
    ],
    """up_block_types""": [
        """AttnUpBlock2D""",
        """AttnUpBlock2D""",
        """AttnUpBlock2D""",
        """ResnetUpsampleBlock2D""",
        """ResnetUpsampleBlock2D""",
        """ResnetUpsampleBlock2D""",
    ],
    """resnet_time_scale_shift""": """default""",
    """upsample_type""": """resnet""",
    """downsample_type""": """resnet""",
}

UpperCAmelCase							 =					{
    """num_train_timesteps""": 40,
    """sigma_min""": 0.0_02,
    """sigma_max""": 80.0,
}

UpperCAmelCase							 =					{
    """num_train_timesteps""": 201,
    """sigma_min""": 0.0_02,
    """sigma_max""": 80.0,
}

UpperCAmelCase							 =					{
    """num_train_timesteps""": 151,
    """sigma_min""": 0.0_02,
    """sigma_max""": 80.0,
}

def 			A_      (	__a     :				Any		):

				"""simple docstring"""
				if isinstance(__a ,		__a		):
								return v
				if v.lower() in ("yes", "true", "t", "y", "1"):
								return True
				elif v.lower() in ("no", "false", "f", "n", "0"):
								return False
				else:
								raise argparse.ArgumentTypeError("""boolean value expected"""		)

def 			A_      (	__a     :				Union[str, Any] ,		__a     :				List[str] ,		__a     :				Dict ,		__a     :				List[str] ,		__a     :				List[str]=False		):

				"""simple docstring"""
				a__				      =		checkpoint[F'''{old_prefix}.in_layers.0.weight''']
				a__				      =		checkpoint[F'''{old_prefix}.in_layers.0.bias''']
				a__				      =		checkpoint[F'''{old_prefix}.in_layers.2.weight''']
				a__				      =		checkpoint[F'''{old_prefix}.in_layers.2.bias''']
				a__				      =		checkpoint[F'''{old_prefix}.emb_layers.1.weight''']
				a__				      =		checkpoint[F'''{old_prefix}.emb_layers.1.bias''']
				a__				      =		checkpoint[F'''{old_prefix}.out_layers.0.weight''']
				a__				      =		checkpoint[F'''{old_prefix}.out_layers.0.bias''']
				a__				      =		checkpoint[F'''{old_prefix}.out_layers.3.weight''']
				a__				      =		checkpoint[F'''{old_prefix}.out_layers.3.bias''']

				if has_skip:
								a__				      =		checkpoint[F'''{old_prefix}.skip_connection.weight''']
								a__				      =		checkpoint[F'''{old_prefix}.skip_connection.bias''']

				return new_checkpoint

def 			A_      (	__a     :				List[Any] ,		__a     :				Optional[int] ,		__a     :				Union[str, Any] ,		__a     :				Dict ,		__a     :				Tuple=None		):

				"""simple docstring"""
				a__      ,			a__      ,			a__				      =		checkpoint[F'''{old_prefix}.qkv.weight'''].chunk(3 ,		dim=0		)
				a__      ,			a__      ,			a__				      =		checkpoint[F'''{old_prefix}.qkv.bias'''].chunk(3 ,		dim=0		)

				a__				      =		checkpoint[F'''{old_prefix}.norm.weight''']
				a__				      =		checkpoint[F'''{old_prefix}.norm.bias''']

				a__				      =		weight_q.squeeze(-1		).squeeze(-1		)
				a__				      =		bias_q.squeeze(-1		).squeeze(-1		)
				a__				      =		weight_k.squeeze(-1		).squeeze(-1		)
				a__				      =		bias_k.squeeze(-1		).squeeze(-1		)
				a__				      =		weight_v.squeeze(-1		).squeeze(-1		)
				a__				      =		bias_v.squeeze(-1		).squeeze(-1		)

				a__				      =		(
				    checkpoint[F'''{old_prefix}.proj_out.weight'''].squeeze(-1		).squeeze(-1		)
				)
				a__				      =		checkpoint[F'''{old_prefix}.proj_out.bias'''].squeeze(-1		).squeeze(-1		)

				return new_checkpoint

def 			A_      (	__a     :				str ,		__a     :				int		):

				"""simple docstring"""
				a__				      =		torch.load(__a ,		map_location="""cpu"""		)
				a__				      =		{}

				a__				      =		checkpoint["""time_embed.0.weight"""]
				a__				      =		checkpoint["""time_embed.0.bias"""]
				a__				      =		checkpoint["""time_embed.2.weight"""]
				a__				      =		checkpoint["""time_embed.2.bias"""]

				if unet_config["num_class_embeds"] is not None:
								a__				      =		checkpoint["""label_emb.weight"""]

				a__				      =		checkpoint["""input_blocks.0.0.weight"""]
				a__				      =		checkpoint["""input_blocks.0.0.bias"""]

				a__				      =		unet_config["""down_block_types"""]
				a__				      =		unet_config["""layers_per_block"""]
				a__				      =		unet_config["""attention_head_dim"""]
				a__				      =		unet_config["""block_out_channels"""]
				a__				      =		1
				a__				      =		channels_list[0]

				for i, layer_type in enumerate(__a		):
								a__				      =		channels_list[i]
								a__				      =		current_channels != prev_channels
								if layer_type == "ResnetDownsampleBlock2D":
												for j in range(__a		):
																a__				      =		F'''down_blocks.{i}.resnets.{j}'''
																a__				      =		F'''input_blocks.{current_layer}.0'''
																a__				      =		True if j == 0 and downsample_block_has_skip else False
																a__				      =		convert_resnet(__a ,		__a ,		__a ,		__a ,		has_skip=__a		)
																current_layer += 1

								elif layer_type == "AttnDownBlock2D":
												for j in range(__a		):
																a__				      =		F'''down_blocks.{i}.resnets.{j}'''
																a__				      =		F'''input_blocks.{current_layer}.0'''
																a__				      =		True if j == 0 and downsample_block_has_skip else False
																a__				      =		convert_resnet(__a ,		__a ,		__a ,		__a ,		has_skip=__a		)
																a__				      =		F'''down_blocks.{i}.attentions.{j}'''
																a__				      =		F'''input_blocks.{current_layer}.1'''
																a__				      =		convert_attention(
																    __a ,		__a ,		__a ,		__a ,		__a		)
																current_layer += 1

								if i != len(__a		) - 1:
												a__				      =		F'''down_blocks.{i}.downsamplers.0'''
												a__				      =		F'''input_blocks.{current_layer}.0'''
												a__				      =		convert_resnet(__a ,		__a ,		__a ,		__a		)
												current_layer += 1

								a__				      =		current_channels

				# hardcoded the mid-block for now
				a__				      =		"""mid_block.resnets.0"""
				a__				      =		"""middle_block.0"""
				a__				      =		convert_resnet(__a ,		__a ,		__a ,		__a		)
				a__				      =		"""mid_block.attentions.0"""
				a__				      =		"""middle_block.1"""
				a__				      =		convert_attention(__a ,		__a ,		__a ,		__a ,		__a		)
				a__				      =		"""mid_block.resnets.1"""
				a__				      =		"""middle_block.2"""
				a__				      =		convert_resnet(__a ,		__a ,		__a ,		__a		)

				a__				      =		0
				a__				      =		unet_config["""up_block_types"""]

				for i, layer_type in enumerate(__a		):
								if layer_type == "ResnetUpsampleBlock2D":
												for j in range(layers_per_block + 1		):
																a__				      =		F'''up_blocks.{i}.resnets.{j}'''
																a__				      =		F'''output_blocks.{current_layer}.0'''
																a__				      =		convert_resnet(__a ,		__a ,		__a ,		__a ,		has_skip=__a		)
																current_layer += 1

												if i != len(__a		) - 1:
																a__				      =		F'''up_blocks.{i}.upsamplers.0'''
																a__				      =		F'''output_blocks.{current_layer-1}.1'''
																a__				      =		convert_resnet(__a ,		__a ,		__a ,		__a		)
								elif layer_type == "AttnUpBlock2D":
												for j in range(layers_per_block + 1		):
																a__				      =		F'''up_blocks.{i}.resnets.{j}'''
																a__				      =		F'''output_blocks.{current_layer}.0'''
																a__				      =		convert_resnet(__a ,		__a ,		__a ,		__a ,		has_skip=__a		)
																a__				      =		F'''up_blocks.{i}.attentions.{j}'''
																a__				      =		F'''output_blocks.{current_layer}.1'''
																a__				      =		convert_attention(
																    __a ,		__a ,		__a ,		__a ,		__a		)
																current_layer += 1

												if i != len(__a		) - 1:
																a__				      =		F'''up_blocks.{i}.upsamplers.0'''
																a__				      =		F'''output_blocks.{current_layer-1}.2'''
																a__				      =		convert_resnet(__a ,		__a ,		__a ,		__a		)

				a__				      =		checkpoint["""out.0.weight"""]
				a__				      =		checkpoint["""out.0.bias"""]
				a__				      =		checkpoint["""out.2.weight"""]
				a__				      =		checkpoint["""out.2.bias"""]

				return new_checkpoint


if __name__ == "__main__":
	UpperCAmelCase							 =					argparse.ArgumentParser()

	parser.add_argument("""--unet_path""", default=None, type=str, required=True, help="""Path to the unet.pt to convert.""")
	parser.add_argument(
	    """--dump_path""", default=None, type=str, required=True, help="""Path to output the converted UNet model."""
	)
	parser.add_argument("""--class_cond""", default=True, type=str, help="""Whether the model is class-conditional.""")

	UpperCAmelCase							 =					parser.parse_args()
	UpperCAmelCase							 =					strabool(args.class_cond)

	UpperCAmelCase							 =					os.path.basename(args.unet_path)
	print(f"""Checkpoint: {ckpt_name}""")

	# Get U-Net config
	if "imagenet64" in ckpt_name:
		UpperCAmelCase							 =					IMAGENET_64_UNET_CONFIG
	elif "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)):
		UpperCAmelCase							 =					LSUN_256_UNET_CONFIG
	elif "test" in ckpt_name:
		UpperCAmelCase							 =					TEST_UNET_CONFIG
	else:
		raise ValueError(f"""Checkpoint type {ckpt_name} is not currently supported.""")

	if not args.class_cond:
		UpperCAmelCase							 =					None

	UpperCAmelCase							 =					con_pt_to_diffuser(args.unet_path, unet_config)

	UpperCAmelCase							 =					UNetaDModel(**unet_config)
	image_unet.load_state_dict(converted_unet_ckpt)

	# Get scheduler config
	if "cd" in ckpt_name or "test" in ckpt_name:
		UpperCAmelCase							 =					CD_SCHEDULER_CONFIG
	elif "ct" in ckpt_name and "imagenet64" in ckpt_name:
		UpperCAmelCase							 =					CT_IMAGENET_64_SCHEDULER_CONFIG
	elif "ct" in ckpt_name and "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)):
		UpperCAmelCase							 =					CT_LSUN_256_SCHEDULER_CONFIG
	else:
		raise ValueError(f"""Checkpoint type {ckpt_name} is not currently supported.""")

	UpperCAmelCase							 =					CMStochasticIterativeScheduler(**scheduler_config)

	UpperCAmelCase							 =					ConsistencyModelPipeline(unet=image_unet, scheduler=cm_scheduler)
	consistency_model.save_pretrained(args.dump_path)






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_      (	__a     :				List[str] ,		__a     :				Optional[int]		):

				"""simple docstring"""
				assert isinstance(__a ,		__a		)
				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_      (	__a     :				Tuple ,		__a     :				Dict ,		__a     :				List[str]		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
								a__				      =		JsonDatasetReader(__a ,		cache_dir=__a ,		keep_in_memory=__a		).read()
				_check_json_dataset(__a ,		__a		)

@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_      (	__a     :				int ,		__a     :				Dict ,		__a     :				List[Any]		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		features.copy() if features else default_expected_features
				a__				      =		(
				    Features({feature: Value(__a		) for feature, dtype in features.items()}		) if features is not None else None
				)
				a__				      =		JsonDatasetReader(__a ,		features=__a ,		cache_dir=__a		).read()
				_check_json_dataset(__a ,		__a		)

@pytest.mark.parametrize(
    """features""" ,		[
        None,
        {"""col_3""": """float64""", """col_1""": """string""", """col_2""": """int64"""},
    ] ,		)
def 			A_      (	__a     :				Tuple ,		__a     :				Tuple ,		__a     :				Tuple		):

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

				"""simple docstring"""
				# jsonl_312_path features are {"col_3": "float64", "col_1": "string", "col_2": "int64"}
				a__				      =		{"""col_2""": """int64""", """col_3""": """float64""", """col_1""": """string"""}
				a__				      =		features.copy()
				a__				      =		(
				    Features({feature: Value(__a		) for feature, dtype in features.items()}		) if features is not None else None
				)
				a__				      =		tmp_path / """cache"""
				a__				      =		JsonDatasetReader(__a ,		features=__a ,		cache_dir=__a		).read()
				assert isinstance(__a ,		__a		)
				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_      (	__a     :				List[Any] ,		__a     :				int ,		__a     :				List[str]		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		JsonDatasetReader(__a ,		cache_dir=__a ,		split=__a		).read()
				_check_json_dataset(__a ,		__a		)
				assert dataset.split == split if split else "train"

@pytest.mark.parametrize("""path_type""" ,		[str, list]		)
def 			A_      (	__a     :				Dict ,		__a     :				Tuple ,		__a     :				List[str]		):

				"""simple docstring"""
				if issubclass(__a ,		__a		):
								a__				      =		jsonl_path
				elif issubclass(__a ,		__a		):
								a__				      =		[jsonl_path]
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		JsonDatasetReader(__a ,		cache_dir=__a		).read()
				_check_json_dataset(__a ,		__a		)

def 			A_      (	__a     :				str ,		__a     :				List[Any] ,		__a     :				Dict=("train",)		):

				"""simple docstring"""
				assert isinstance(__a ,		__a		)
				for split in splits:
								a__				      =		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_      (	__a     :				int ,		__a     :				int ,		__a     :				Dict		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
								a__				      =		JsonDatasetReader({"""train""": jsonl_path} ,		cache_dir=__a ,		keep_in_memory=__a		).read()
				_check_json_datasetdict(__a ,		__a		)

@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_      (	__a     :				Tuple ,		__a     :				Dict ,		__a     :				Optional[Any]		):

				"""simple docstring"""
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		features.copy() if features else default_expected_features
				a__				      =		(
				    Features({feature: Value(__a		) for feature, dtype in features.items()}		) if features is not None else None
				)
				a__				      =		JsonDatasetReader({"""train""": jsonl_path} ,		features=__a ,		cache_dir=__a		).read()
				_check_json_datasetdict(__a ,		__a		)

@pytest.mark.parametrize("""split""" ,		[None, NamedSplit("""train"""		), """train""", """test"""]		)
def 			A_      (	__a     :				Dict ,		__a     :				Any ,		__a     :				Tuple		):

				"""simple docstring"""
				if split:
								a__				      =		{split: jsonl_path}
				else:
								a__				      =		"""train"""
								a__				      =		{"""train""": jsonl_path, """test""": jsonl_path}
				a__				      =		tmp_path / """cache"""
				a__				      =		{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
				a__				      =		JsonDatasetReader(__a ,		cache_dir=__a		).read()
				_check_json_datasetdict(__a ,		__a ,		splits=list(path.keys()		)		)
				assert all(dataset[split].split == split for split in path.keys()		)

def 			A_      (	__a     :				Any		):

				"""simple docstring"""
				return json.load(__a		)

def 			A_      (	__a     :				Tuple		):

				"""simple docstring"""
				return [json.loads(__a		) for line in buffer]




class  __snake_case      :







							'''simple docstring'''






							@pytest.mark.parametrize("""lines, load_json_function"""  ,  [(True, load_json_lines), (False, load_json)]		)
							def _a      (							self  ,  a_  ,  a_  ,  a_		):
											with io.BytesIO() as buffer:
															JsonDatasetWriter(a_  ,  a_  ,  lines=a_		).write()
															buffer.seek(0		)
															a__				      =		load_json_function(a_		)
											assert isinstance(a_  ,  a_		)
											assert isinstance(exported_content[0]  ,  a_		)
											assert len(a_		) == 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 _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											with io.BytesIO() as buffer:
															JsonDatasetWriter(a_  ,  a_  ,  lines=a_  ,  orient=a_		).write()
															buffer.seek(0		)
															a__				      =		load_json(a_		)
											assert isinstance(a_  ,  a_		)
											if keys:
															if container is dict:
																			assert exported_content.keys() == keys
															else:
																			assert exported_content[0].keys() == keys
											else:
															assert not hasattr(a_  ,  """keys"""		) and not hasattr(exported_content[0]  ,  """keys"""		)
											if len_at:
															assert len(exported_content[len_at]		) == 10
											else:
															assert len(a_		) == 10


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


							def _a      (							self  ,  a_		):
											with pytest.raises(a_		):
															with io.BytesIO() as buffer:
																			JsonDatasetWriter(a_  ,  a_  ,  num_proc=0		)







							@pytest.mark.parametrize("""compression, extension"""  ,  [("""gzip""", """gz"""), ("""bz2""", """bz2"""), ("""xz""", """xz""")]		)
							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		tmp_path_factory.mktemp("""data"""		) / F'''test.json.{extension}'''
											a__				      =		str(shared_datadir / F'''test_file.json.{extension}'''		)
											JsonDatasetWriter(a_  ,  a_  ,  compression=a_		).write()

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

