



import argparse
import os
import re

import tensorflow as tf
import torch

from transformers import BertConfig, BertModel
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase	:		int				=   logging.get_logger(__name__)



def   __lowerCAmelCase					(							lowerCamelCase   :						Union[str, Any]				,		lowerCamelCase   :						Optional[Any]				,		lowerCamelCase   :						Union[str, Any]					):






	'''simple docstring'''




	__lowerCAmelCase												=  os.path.abspath(lowerCamelCase_					)
	logger.info(f'''Converting TensorFlow checkpoint from {tf_path}'''					)
	# Load weights from TF model
	__lowerCAmelCase												=  tf.train.list_variables(lowerCamelCase_					)
	__lowerCAmelCase												=  []
	__lowerCAmelCase												=  []
	__lowerCAmelCase												=  []
	for full_name, shape in init_vars:
		# logger.info(f"Loading TF weight {name} with shape {shape}")
		__lowerCAmelCase												=  full_name.split("/"					)
		if full_name == "_CHECKPOINTABLE_OBJECT_GRAPH" or name[0] in ["global_step", "save_counter"]:
			logger.info(f'''Skipping non-model layer {full_name}'''					)
			continue
		if "optimizer" in full_name:
			logger.info(f'''Skipping optimization layer {full_name}'''					)
			continue
		if name[0] == "model":
			# ignore initial 'model'
			__lowerCAmelCase												=  name[1:]
		# figure out how many levels deep the name is
		__lowerCAmelCase												=  0
		for _name in name:
			if _name.startswith("layer_with_weights"					):
				depth += 1
			else:
				break
		layer_depth.append(lowerCamelCase_					)
		# read data
		__lowerCAmelCase												=  tf.train.load_variable(lowerCamelCase_				,		lowerCamelCase_					)
		names.append("/".join(lowerCamelCase_					)					)
		arrays.append(lowerCamelCase_					)
	logger.info(f'''Read a total of {len(lowerCamelCase_					):,} layers'''					)

	# Sanity check
	if len(set(lowerCamelCase_					)					) != 1:
		raise ValueError(f'''Found layer names with different depths (layer depth {list(set(lowerCamelCase_					)					)})'''					)
	__lowerCAmelCase												=  list(set(lowerCamelCase_					)					)[0]
	if layer_depth != 1:
		raise ValueError(
		    "The model contains more than just the embedding/encoder layers. This script does not handle MLM/NSP"
		    " heads."					)

	# convert layers
	logger.info("Converting weights..."					)
	for full_name, array in zip(lowerCamelCase_				,		lowerCamelCase_					):
		__lowerCAmelCase												=  full_name.split("/"					)
		__lowerCAmelCase												=  model
		__lowerCAmelCase												=  []
		for i, m_name in enumerate(lowerCamelCase_					):
			if m_name == ".ATTRIBUTES":
				# variable names end with .ATTRIBUTES/VARIABLE_VALUE
				break
			if m_name.startswith("layer_with_weights"					):
				__lowerCAmelCase												=  int(m_name.split("-"					)[-1]					)
				if layer_num <= 2:
					# embedding layers
					# layer_num 0: word_embeddings
					# layer_num 1: position_embeddings
					# layer_num 2: token_type_embeddings
					continue
				elif layer_num == 3:
					# embedding LayerNorm
					trace.extend(["embeddings", "LayerNorm"]					)
					__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"embeddings"					)
					__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"LayerNorm"					)
				elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
					# encoder layers
					trace.extend(["encoder", "layer", str(layer_num - 4					)]					)
					__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"encoder"					)
					__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"layer"					)
					__lowerCAmelCase												=  pointer[layer_num - 4]
				elif layer_num == config.num_hidden_layers + 4:
					# pooler layer
					trace.extend(["pooler", "dense"]					)
					__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"pooler"					)
					__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"dense"					)
			elif m_name == "embeddings":
				trace.append("embeddings"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"embeddings"					)
				if layer_num == 0:
					trace.append("word_embeddings"					)
					__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"word_embeddings"					)
				elif layer_num == 1:
					trace.append("position_embeddings"					)
					__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"position_embeddings"					)
				elif layer_num == 2:
					trace.append("token_type_embeddings"					)
					__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"token_type_embeddings"					)
				else:
					raise ValueError(f'''Unknown embedding layer with name {full_name}'''					)
				trace.append("weight"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"weight"					)
			elif m_name == "_attention_layer":
				# self-attention layer
				trace.extend(["attention", "self"]					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"attention"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"self"					)
			elif m_name == "_attention_layer_norm":
				# output attention norm
				trace.extend(["attention", "output", "LayerNorm"]					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"attention"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"output"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"LayerNorm"					)
			elif m_name == "_attention_output_dense":
				# output attention dense
				trace.extend(["attention", "output", "dense"]					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"attention"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"output"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"dense"					)
			elif m_name == "_output_dense":
				# output dense
				trace.extend(["output", "dense"]					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"output"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"dense"					)
			elif m_name == "_output_layer_norm":
				# output dense
				trace.extend(["output", "LayerNorm"]					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"output"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"LayerNorm"					)
			elif m_name == "_key_dense":
				# attention key
				trace.append("key"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"key"					)
			elif m_name == "_query_dense":
				# attention query
				trace.append("query"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"query"					)
			elif m_name == "_value_dense":
				# attention value
				trace.append("value"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"value"					)
			elif m_name == "_intermediate_dense":
				# attention intermediate dense
				trace.extend(["intermediate", "dense"]					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"intermediate"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"dense"					)
			elif m_name == "_output_layer_norm":
				# output layer norm
				trace.append("output"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"output"					)
			# weights & biases
			elif m_name in ["bias", "beta"]:
				trace.append("bias"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"bias"					)
			elif m_name in ["kernel", "gamma"]:
				trace.append("weight"					)
				__lowerCAmelCase												=  getattr(lowerCamelCase_				,		"weight"					)
			else:
				logger.warning(f'''Ignored {m_name}'''					)
        # for certain layers reshape is necessary
		__lowerCAmelCase												=  '''.'''.join(lowerCamelCase_					)
		if re.match(r"(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)"				,		lowerCamelCase_					) or re.match(
		    r"(\S+)\.attention\.output\.dense\.weight"				,		lowerCamelCase_					):
			__lowerCAmelCase												=  array.reshape(pointer.data.shape					)
		if "kernel" in full_name:
			__lowerCAmelCase												=  array.transpose()
		if pointer.shape == array.shape:
			__lowerCAmelCase												=  torch.from_numpy(lowerCamelCase_					)
		else:
			raise ValueError(
			    f'''Shape mismatch in layer {full_name}: Model expects shape {pointer.shape} but layer contains shape:'''
			    f''' {array.shape}'''					)
		logger.info(f'''Successfully set variable {full_name} to PyTorch layer {trace}'''					)
	return model



def   __lowerCAmelCase					(							lowerCamelCase   :						Optional[Any]				,		lowerCamelCase   :						str				,		lowerCamelCase   :						List[str]					):






	'''simple docstring'''




	logger.info(f'''Loading model based on config from {config_path}...'''					)
	__lowerCAmelCase												=  BertConfig.from_json_file(lowerCamelCase_					)
	__lowerCAmelCase												=  BertModel(lowerCamelCase_					)

	# Load weights from checkpoint
	logger.info(f'''Loading weights from checkpoint {tf_checkpoint_path}...'''					)
	load_tfa_weights_in_bert(lowerCamelCase_				,		lowerCamelCase_				,		lowerCamelCase_					)

	# Save pytorch-model
	logger.info(f'''Saving PyTorch model to {pytorch_dump_path}...'''					)
	torch.save(model.state_dict()				,		lowerCamelCase_					)


if __name__ == "__main__":
	lowerCAmelCase	:		int				=   argparse.ArgumentParser()
	parser.add_argument(
	    '''--tf_checkpoint_path''', type=str, required=True, help='''Path to the TensorFlow 2.x checkpoint path.'''
	)
	parser.add_argument(
	    '''--bert_config_file''',
	    type=str,
	    required=True,
	    help='''The config json file corresponding to the BERT model. This specifies the model architecture.''',
	)
	parser.add_argument(
	    '''--pytorch_dump_path''',
	    type=str,
	    required=True,
	    help='''Path to the output PyTorch model (must include filename).''',
	)
	lowerCAmelCase	:		Tuple				=   parser.parse_args()
	convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_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   __lowerCAmelCase					(							lowerCamelCase   :						str				,		lowerCamelCase   :						Union[str, Any]					):






	'''simple docstring'''




	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   :						Tuple				,		lowerCamelCase   :						int				,		lowerCamelCase   :						Optional[int]					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
		__lowerCAmelCase												=  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   :						List[Any]				,		lowerCamelCase   :						int				,		lowerCamelCase   :						Dict					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	__lowerCAmelCase												=  features.copy() if features else default_expected_features
	__lowerCAmelCase												=  (
	    Features({feature: Value(lowerCamelCase					) for feature, dtype in features.items()}					) if features is not None else None
	)
	__lowerCAmelCase												=  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   :						List[Any]				,		lowerCamelCase   :						Optional[int]				,		lowerCamelCase   :						Any					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_3": "float64", "col_1": "string", "col_2": "int64"}
	__lowerCAmelCase												=  features.copy() if features else default_expected_features
	__lowerCAmelCase												=  (
	    Features({feature: Value(lowerCamelCase					) for feature, dtype in features.items()}					) if features is not None else None
	)
	__lowerCAmelCase												=  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   :						Any				,		lowerCamelCase   :						Optional[int]					):






	'''simple docstring'''




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






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	__lowerCAmelCase												=  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   :						Any				,		lowerCamelCase   :						List[str]				,		lowerCamelCase   :						Optional[int]					):






	'''simple docstring'''




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



def   __lowerCAmelCase					(							lowerCamelCase   :						Optional[int]				,		lowerCamelCase   :						int				,		lowerCamelCase   :						str=("train",)					):






	'''simple docstring'''




	assert isinstance(lowerCamelCase				,		lowerCamelCase					)
	for split in splits:
		__lowerCAmelCase												=  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   :						Dict				,		lowerCamelCase   :						int				,		lowerCamelCase   :						List[str]					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
		__lowerCAmelCase												=  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   :						int				,		lowerCamelCase   :						Any				,		lowerCamelCase   :						Dict					):






	'''simple docstring'''




	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	__lowerCAmelCase												=  features.copy() if features else default_expected_features
	__lowerCAmelCase												=  (
	    Features({feature: Value(lowerCamelCase					) for feature, dtype in features.items()}					) if features is not None else None
	)
	__lowerCAmelCase												=  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   :						int				,		lowerCamelCase   :						Optional[int]				,		lowerCamelCase   :						int					):






	'''simple docstring'''




	if split:
		__lowerCAmelCase												=  {split: jsonl_path}
	else:
		__lowerCAmelCase												=  "train"
		__lowerCAmelCase												=  {"train": jsonl_path, "test": jsonl_path}
	__lowerCAmelCase												=  tmp_path / "cache"
	__lowerCAmelCase												=  {"col_1": "string", "col_2": "int64", "col_3": "float64"}
	__lowerCAmelCase												=  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   :						Optional[Any]					):






	'''simple docstring'''




	return json.load(lowerCamelCase					)



def   __lowerCAmelCase					(							lowerCamelCase   :						Tuple					):






	'''simple docstring'''




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






class 					UpperCAmelCase__	:



					@pytest.mark.parametrize("lines, load_json_function"			,				[(True, load_json_lines), (False, load_json)]     )
					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			int:
						with io.BytesIO() as buffer:
							JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				lines=UpperCamelCase     ).write()
							buffer.seek(0     )
							__lowerCAmelCase												=  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        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			int:
						with io.BytesIO() as buffer:
							JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				lines=UpperCamelCase			,				orient=UpperCamelCase     ).write()
							buffer.seek(0     )
							__lowerCAmelCase												=  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        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			str:
						with io.BytesIO() as buffer:
							JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				lines=UpperCamelCase			,				num_proc=2     ).write()
							buffer.seek(0     )
							__lowerCAmelCase												=  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        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			Union[str, Any]:
						with io.BytesIO() as buffer:
							JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				lines=UpperCamelCase			,				orient=UpperCamelCase			,				num_proc=2     ).write()
							buffer.seek(0     )
							__lowerCAmelCase												=  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        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			Any:
						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        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			Tuple:
						__lowerCAmelCase												=  tmp_path_factory.mktemp("data"     ) / F'''test.json.{extension}'''
						__lowerCAmelCase												=  str(shared_datadir / F'''test_file.json.{extension}'''     )
						JsonDatasetWriter(UpperCamelCase			,				UpperCamelCase			,				compression=UpperCamelCase     ).write()

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