'''simple docstring'''





import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


# =================#
# UNet Conversion #
# =================#

SCREAMING_SNAKE_CASE_:	Dict    =[
    # (stable-diffusion, HF Diffusers)
    ('time_embed.0.weight', 'time_embedding.linear_1.weight'),
    ('time_embed.0.bias', 'time_embedding.linear_1.bias'),
    ('time_embed.2.weight', 'time_embedding.linear_2.weight'),
    ('time_embed.2.bias', 'time_embedding.linear_2.bias'),
    ('input_blocks.0.0.weight', 'conv_in.weight'),
    ('input_blocks.0.0.bias', 'conv_in.bias'),
    ('out.0.weight', 'conv_norm_out.weight'),
    ('out.0.bias', 'conv_norm_out.bias'),
    ('out.2.weight', 'conv_out.weight'),
    ('out.2.bias', 'conv_out.bias'),
]

SCREAMING_SNAKE_CASE_:	List[Any]    =[
    # (stable-diffusion, HF Diffusers)
    ('in_layers.0', 'norm1'),
    ('in_layers.2', 'conv1'),
    ('out_layers.0', 'norm2'),
    ('out_layers.3', 'conv2'),
    ('emb_layers.1', 'time_emb_proj'),
    ('skip_connection', 'conv_shortcut'),
]

SCREAMING_SNAKE_CASE_:	Union[str, Any]    =[]
# hardcoded number of downblocks and resnets/attentions...
# would need smarter logic for other networks.
for i in range(4):
    # loop over downblocks/upblocks

    for j in range(2):
        # loop over resnets/attentions for downblocks
        SCREAMING_SNAKE_CASE_:	Any    =f"down_blocks.{i}.resnets.{j}."
        SCREAMING_SNAKE_CASE_:	Tuple    =f"input_blocks.{3*i + j + 1}.0."
        unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))

        if i < 3:
            # no attention layers in down_blocks.3
            SCREAMING_SNAKE_CASE_:	Optional[Any]    =f"down_blocks.{i}.attentions.{j}."
            SCREAMING_SNAKE_CASE_:	List[str]    =f"input_blocks.{3*i + j + 1}.1."
            unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))

    for j in range(3):
        # loop over resnets/attentions for upblocks
        SCREAMING_SNAKE_CASE_:	Union[str, Any]    =f"up_blocks.{i}.resnets.{j}."
        SCREAMING_SNAKE_CASE_:	Any    =f"output_blocks.{3*i + j}.0."
        unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))

        if i > 0:
            # no attention layers in up_blocks.0
            SCREAMING_SNAKE_CASE_:	int    =f"up_blocks.{i}.attentions.{j}."
            SCREAMING_SNAKE_CASE_:	Optional[int]    =f"output_blocks.{3*i + j}.1."
            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))

    if i < 3:
        # no downsample in down_blocks.3
        SCREAMING_SNAKE_CASE_:	Union[str, Any]    =f"down_blocks.{i}.downsamplers.0.conv."
        SCREAMING_SNAKE_CASE_:	Union[str, Any]    =f"input_blocks.{3*(i+1)}.0.op."
        unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))

        # no upsample in up_blocks.3
        SCREAMING_SNAKE_CASE_:	int    =f"up_blocks.{i}.upsamplers.0."
        SCREAMING_SNAKE_CASE_:	List[Any]    =f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
        unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

SCREAMING_SNAKE_CASE_:	int    ='mid_block.attentions.0.'
SCREAMING_SNAKE_CASE_:	List[Any]    ='middle_block.1.'
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
    SCREAMING_SNAKE_CASE_:	Tuple    =f"mid_block.resnets.{j}."
    SCREAMING_SNAKE_CASE_:	Tuple    =f"middle_block.{2*j}."
    unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))





def       lowerCAmelCase_							(   snake_case_			:					Optional[Any]     )       -> List[str]:


     '''simple docstring'''




     UpperCAmelCase_			       =  {k: k for k in unet_state_dict.keys()}
     for sd_name, hf_name in unet_conversion_map:
          UpperCAmelCase_			       =  sd_name
     for k, v in mapping.items():
          if "resnets" in k:
               for sd_part, hf_part in unet_conversion_map_resnet:
                    UpperCAmelCase_			       =  v.replace(snake_case_   ,       snake_case_     )
               UpperCAmelCase_			       =  v
     for k, v in mapping.items():
          for sd_part, hf_part in unet_conversion_map_layer:
               UpperCAmelCase_			       =  v.replace(snake_case_   ,       snake_case_     )
          UpperCAmelCase_			       =  v
     UpperCAmelCase_			       =  {v: unet_state_dict[k] for k, v in mapping.items()}
     return new_state_dict


# ================#
# VAE Conversion #
# ================#

SCREAMING_SNAKE_CASE_:	int    =[
    # (stable-diffusion, HF Diffusers)
    ('nin_shortcut', 'conv_shortcut'),
    ('norm_out', 'conv_norm_out'),
    ('mid.attn_1.', 'mid_block.attentions.0.'),
]

for i in range(4):
    # down_blocks have two resnets
    for j in range(2):
        SCREAMING_SNAKE_CASE_:	Tuple    =f"encoder.down_blocks.{i}.resnets.{j}."
        SCREAMING_SNAKE_CASE_:	int    =f"encoder.down.{i}.block.{j}."
        vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

    if i < 3:
        SCREAMING_SNAKE_CASE_:	int    =f"down_blocks.{i}.downsamplers.0."
        SCREAMING_SNAKE_CASE_:	str    =f"down.{i}.downsample."
        vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

        SCREAMING_SNAKE_CASE_:	int    =f"up_blocks.{i}.upsamplers.0."
        SCREAMING_SNAKE_CASE_:	List[str]    =f"up.{3-i}.upsample."
        vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))

    # up_blocks have three resnets
    # also, up blocks in hf are numbered in reverse from sd
    for j in range(3):
        SCREAMING_SNAKE_CASE_:	List[str]    =f"decoder.up_blocks.{i}.resnets.{j}."
        SCREAMING_SNAKE_CASE_:	Dict    =f"decoder.up.{3-i}.block.{j}."
        vae_conversion_map.append((sd_up_prefix, hf_up_prefix))

# this part accounts for mid blocks in both the encoder and the decoder
for i in range(2):
    SCREAMING_SNAKE_CASE_:	Any    =f"mid_block.resnets.{i}."
    SCREAMING_SNAKE_CASE_:	Tuple    =f"mid.block_{i+1}."
    vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


SCREAMING_SNAKE_CASE_:	int    =[
    # (stable-diffusion, HF Diffusers)
    ('norm.', 'group_norm.'),
    ('q.', 'query.'),
    ('k.', 'key.'),
    ('v.', 'value.'),
    ('proj_out.', 'proj_attn.'),
]





def       lowerCAmelCase_							(   snake_case_			:					Tuple     )       -> Tuple:


     '''simple docstring'''




     return w.reshape(*w.shape   ,       1   ,       1     )





def       lowerCAmelCase_							(   snake_case_			:					Optional[Any]     )       -> Optional[Any]:


     '''simple docstring'''




     UpperCAmelCase_			       =  {k: k for k in vae_state_dict.keys()}
     for k, v in mapping.items():
          for sd_part, hf_part in vae_conversion_map:
               UpperCAmelCase_			       =  v.replace(snake_case_   ,       snake_case_     )
          UpperCAmelCase_			       =  v
     for k, v in mapping.items():
          if "attentions" in k:
               for sd_part, hf_part in vae_conversion_map_attn:
                    UpperCAmelCase_			       =  v.replace(snake_case_   ,       snake_case_     )
               UpperCAmelCase_			       =  v
     UpperCAmelCase_			       =  {v: vae_state_dict[k] for k, v in mapping.items()}
     UpperCAmelCase_			       =  ["q", "k", "v", "proj_out"]
     for k, v in new_state_dict.items():
          for weight_name in weights_to_convert:
               if f"""mid.attn_1.{weight_name}.weight""" in k:
                    print(f"""Reshaping {k} for SD format"""     )
                    UpperCAmelCase_			       =  reshape_weight_for_sd(snake_case_     )
     return new_state_dict


# =========================#
# Text Encoder Conversion #
# =========================#


SCREAMING_SNAKE_CASE_:	List[Any]    =[
    # (stable-diffusion, HF Diffusers)
    ('resblocks.', 'text_model.encoder.layers.'),
    ('ln_1', 'layer_norm1'),
    ('ln_2', 'layer_norm2'),
    ('.c_fc.', '.fc1.'),
    ('.c_proj.', '.fc2.'),
    ('.attn', '.self_attn'),
    ('ln_final.', 'transformer.text_model.final_layer_norm.'),
    ('token_embedding.weight', 'transformer.text_model.embeddings.token_embedding.weight'),
    ('positional_embedding', 'transformer.text_model.embeddings.position_embedding.weight'),
]
SCREAMING_SNAKE_CASE_:	Dict    ={re.escape(x[1]): x[0] for x in textenc_conversion_lst}
SCREAMING_SNAKE_CASE_:	str    =re.compile('|'.join(protected.keys()))

# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
SCREAMING_SNAKE_CASE_:	List[Any]    ={'q': 0, 'k': 1, 'v': 2}





def       lowerCAmelCase_							(   snake_case_			:					Union[str, Any]     )       -> Tuple:


     '''simple docstring'''




     UpperCAmelCase_			       =  {}
     UpperCAmelCase_			       =  {}
     UpperCAmelCase_			       =  {}
     for k, v in text_enc_dict.items():
          if (
              k.endswith(".self_attn.q_proj.weight"     )
              or k.endswith(".self_attn.k_proj.weight"     )
              or k.endswith(".self_attn.v_proj.weight"     )
          ):
               UpperCAmelCase_			       =  k[: -len(".q_proj.weight"     )]
               UpperCAmelCase_			       =  k[-len("q_proj.weight"     )]
               if k_pre not in capture_qkv_weight:
                    UpperCAmelCase_			       =  [None, None, None]
               UpperCAmelCase_			       =  v
               continue

          if (
              k.endswith(".self_attn.q_proj.bias"     )
              or k.endswith(".self_attn.k_proj.bias"     )
              or k.endswith(".self_attn.v_proj.bias"     )
          ):
               UpperCAmelCase_			       =  k[: -len(".q_proj.bias"     )]
               UpperCAmelCase_			       =  k[-len("q_proj.bias"     )]
               if k_pre not in capture_qkv_bias:
                    UpperCAmelCase_			       =  [None, None, None]
               UpperCAmelCase_			       =  v
               continue

          UpperCAmelCase_			       =  textenc_pattern.sub(lambda snake_case_     : protected[re.escape(m.group(0     )     )]   ,       snake_case_     )
          UpperCAmelCase_			       =  v

     for k_pre, tensors in capture_qkv_weight.items():
          if None in tensors:
               raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing"     )
          UpperCAmelCase_			       =  textenc_pattern.sub(lambda snake_case_     : protected[re.escape(m.group(0     )     )]   ,       snake_case_     )
          UpperCAmelCase_			       =  torch.cat(snake_case_     )

     for k_pre, tensors in capture_qkv_bias.items():
          if None in tensors:
               raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing"     )
          UpperCAmelCase_			       =  textenc_pattern.sub(lambda snake_case_     : protected[re.escape(m.group(0     )     )]   ,       snake_case_     )
          UpperCAmelCase_			       =  torch.cat(snake_case_     )

     return new_state_dict





def       lowerCAmelCase_							(   snake_case_			:					List[Any]     )       -> Union[str, Any]:


     '''simple docstring'''




     return text_enc_dict


if __name__ == "__main__":
    SCREAMING_SNAKE_CASE_:	str    =argparse.ArgumentParser()

    parser.add_argument('--model_path', default=None, type=str, required=True, help='Path to the model to convert.')
    parser.add_argument('--checkpoint_path', default=None, type=str, required=True, help='Path to the output model.')
    parser.add_argument('--half', action='store_true', help='Save weights in half precision.')
    parser.add_argument(
        '--use_safetensors', action='store_true', help='Save weights use safetensors, default is ckpt.'
    )

    SCREAMING_SNAKE_CASE_:	Dict    =parser.parse_args()

    assert args.model_path is not None, "Must provide a model path!"

    assert args.checkpoint_path is not None, "Must provide a checkpoint path!"

    # Path for safetensors
    SCREAMING_SNAKE_CASE_:	Any    =osp.join(args.model_path, 'unet', 'diffusion_pytorch_model.safetensors')
    SCREAMING_SNAKE_CASE_:	Dict    =osp.join(args.model_path, 'vae', 'diffusion_pytorch_model.safetensors')
    SCREAMING_SNAKE_CASE_:	Union[str, Any]    =osp.join(args.model_path, 'text_encoder', 'model.safetensors')

    # Load models from safetensors if it exists, if it doesn't pytorch
    if osp.exists(unet_path):
        SCREAMING_SNAKE_CASE_:	Union[str, Any]    =load_file(unet_path, device='cpu')
    else:
        SCREAMING_SNAKE_CASE_:	int    =osp.join(args.model_path, 'unet', 'diffusion_pytorch_model.bin')
        SCREAMING_SNAKE_CASE_:	Dict    =torch.load(unet_path, map_location='cpu')

    if osp.exists(vae_path):
        SCREAMING_SNAKE_CASE_:	Tuple    =load_file(vae_path, device='cpu')
    else:
        SCREAMING_SNAKE_CASE_:	List[Any]    =osp.join(args.model_path, 'vae', 'diffusion_pytorch_model.bin')
        SCREAMING_SNAKE_CASE_:	str    =torch.load(vae_path, map_location='cpu')

    if osp.exists(text_enc_path):
        SCREAMING_SNAKE_CASE_:	Tuple    =load_file(text_enc_path, device='cpu')
    else:
        SCREAMING_SNAKE_CASE_:	List[Any]    =osp.join(args.model_path, 'text_encoder', 'pytorch_model.bin')
        SCREAMING_SNAKE_CASE_:	Any    =torch.load(text_enc_path, map_location='cpu')

    # Convert the UNet model
    SCREAMING_SNAKE_CASE_:	List[Any]    =convert_unet_state_dict(unet_state_dict)
    SCREAMING_SNAKE_CASE_:	Any    ={'model.diffusion_model.' + k: v for k, v in unet_state_dict.items()}

    # Convert the VAE model
    SCREAMING_SNAKE_CASE_:	List[Any]    =convert_vae_state_dict(vae_state_dict)
    SCREAMING_SNAKE_CASE_:	Dict    ={'first_stage_model.' + k: v for k, v in vae_state_dict.items()}

    # Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
    SCREAMING_SNAKE_CASE_:	Dict    ='text_model.encoder.layers.22.layer_norm2.bias' in text_enc_dict

    if is_vaa_model:
        # Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
        SCREAMING_SNAKE_CASE_:	Any    ={'transformer.' + k: v for k, v in text_enc_dict.items()}
        SCREAMING_SNAKE_CASE_:	str    =convert_text_enc_state_dict_vaa(text_enc_dict)
        SCREAMING_SNAKE_CASE_:	int    ={'cond_stage_model.model.' + k: v for k, v in text_enc_dict.items()}
    else:
        SCREAMING_SNAKE_CASE_:	str    =convert_text_enc_state_dict(text_enc_dict)
        SCREAMING_SNAKE_CASE_:	Optional[int]    ={'cond_stage_model.transformer.' + k: v for k, v in text_enc_dict.items()}

    # Put together new checkpoint
    SCREAMING_SNAKE_CASE_:	List[str]    ={**unet_state_dict, **vae_state_dict, **text_enc_dict}
    if args.half:
        SCREAMING_SNAKE_CASE_:	List[str]    ={k: v.half() for k, v in state_dict.items()}

    if args.use_safetensors:
        save_file(state_dict, args.checkpoint_path)
    else:
        SCREAMING_SNAKE_CASE_:	str    ={'state_dict': state_dict}
        torch.save(state_dict, args.checkpoint_path)











"""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      UpperCamelCase							(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE     ):
						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      UpperCamelCase							(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE     ):
						UpperCamelCase				:		Tuple				   =			tmp_path / """cache"""
						UpperCamelCase				:		Optional[int]				   =			{"""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				:		Tuple				   =			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      UpperCamelCase							(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE     ):
						UpperCamelCase				:		Optional[Any]				   =			tmp_path / """cache"""
						UpperCamelCase				:		int				   =			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
						UpperCamelCase				:		Union[str, Any]				   =			features.copy() if features else default_expected_features
						UpperCamelCase				:		str				   =			(
						    Features({feature: Value(SCREAMING_SNAKE_CASE     ) for feature, dtype in features.items()}     ) if features is not None else None
						)
						UpperCamelCase				:		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      UpperCamelCase							(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE     ):
						UpperCamelCase				:		int				   =			tmp_path / """cache"""
						UpperCamelCase				:		Optional[int]				   =			{"""col_3""": """float64""", """col_1""": """string""", """col_2""": """int64"""}
						UpperCamelCase				:		Optional[Any]				   =			features.copy() if features else default_expected_features
						UpperCamelCase				:		Tuple				   =			(
						    Features({feature: Value(SCREAMING_SNAKE_CASE     ) for feature, dtype in features.items()}     ) if features is not None else None
						)
						UpperCamelCase				:		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_3", "col_1", "col_2"]
						for feature, expected_dtype in expected_features.items():
												assert dataset.features[feature].dtype == expected_dtype

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

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

def      UpperCamelCase							(SCREAMING_SNAKE_CASE     ):
						return [json.loads(SCREAMING_SNAKE_CASE     ) for line in buffer]






class 							lowercase__			:

			"""simple docstring"""







			@pytest.mark.parametrize("""lines, load_json_function"""						,   [(True, load_json_lines), (False, load_json)]    )
			def   _a      (      self						,   _A						,   _A						,   _A    ):






									'''simple docstring'''







									with io.BytesIO() as buffer:
															JsonDatasetWriter(_A						,   _A						,   lines=_A    ).write()
															buffer.seek(0    )
															UpperCamelCase				:		Optional[Any]				   =			load_json_function(_A    )
									assert isinstance(_A						,   _A    )
									assert isinstance(exported_content[0]						,   _A    )
									assert len(_A    ) == 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						,   _A						,   _A						,   _A						,   _A						,   _A    ):






									'''simple docstring'''







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







			@pytest.mark.parametrize("""lines, load_json_function"""						,   [(True, load_json_lines), (False, load_json)]    )
			def   _a      (      self						,   _A						,   _A						,   _A    ):






									'''simple docstring'''







									with io.BytesIO() as buffer:
															JsonDatasetWriter(_A						,   _A						,   lines=_A						,   num_proc=2    ).write()
															buffer.seek(0    )
															UpperCamelCase				:		List[Any]				   =			load_json_function(_A    )
									assert isinstance(_A						,   _A    )
									assert isinstance(exported_content[0]						,   _A    )
									assert len(_A    ) == 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						,   _A						,   _A						,   _A						,   _A						,   _A    ):






									'''simple docstring'''







									with io.BytesIO() as buffer:
															JsonDatasetWriter(_A						,   _A						,   lines=_A						,   orient=_A						,   num_proc=2    ).write()
															buffer.seek(0    )
															UpperCamelCase				:		Dict				   =			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]    ) == 1_0
									else:
															assert len(_A    ) == 1_0







			def   _a      (      self						,   _A    ):






									'''simple docstring'''







									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    ):






									'''simple docstring'''







									UpperCamelCase				:		List[str]				   =			tmp_path_factory.mktemp("""data"""    ) / f"""test.json.{extension}"""
									UpperCamelCase				:		Tuple				   =			str(shared_datadir / f"""test_file.json.{extension}"""    )
									JsonDatasetWriter(_A						,   _A						,   compression=_A    ).write()

									with fsspec.open(_A						,   """rb"""						,   compression="""infer"""    ) as f:
															UpperCamelCase				:		Dict				   =			f.read()
									with fsspec.open(_A						,   """rb"""						,   compression="""infer"""    ) as f:
															UpperCamelCase				:		Tuple				   =			f.read()
									assert exported_content == original_content



